content: # Copyright 2015 Google Inc. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """Tests for yapf.comment_splicer."""
  
  import textwrap
  import unittest
  
  from yapf.pytree import comment_splicer
  from yapf.pytree import pytree_utils
  
  from yapftests import yapf_test_helper
  
  
  class CommentSplicerTest(yapf_test_helper.YAPFTest):
  
    def _AssertNodeType(self, expected_type, node):
      self.assertEqual(expected_type, pytree_utils.NodeName(node))
  
    def _AssertNodeIsComment(self, node, text_in_comment=None):
      if pytree_utils.NodeName(node) == 'simple_stmt':
        self._AssertNodeType('COMMENT', node.children[0])
        node_value = node.children[0].value
      else:
        self._AssertNodeType('COMMENT', node)
        node_value = node.value
      if text_in_comment is not None:
        self.assertIn(text_in_comment, node_value)
  
    def _FindNthChildNamed(self, node, name, n=1):
      for i, child in enumerate(
          [c for c in node.pre_order() if pytree_utils.NodeName(c) == name]):
        if i == n - 1:
          return child
      raise RuntimeError('No Nth child for n={0}'.format(n))
  
    def testSimpleInline(self):
      code = textwrap.dedent("""\
          foo = 1 # and a comment
      """)
      tree = pytree_utils.ParseCodeToTree(code)
      comment_splicer.SpliceComments(tree)
  
      expr = tree.children[0].children[0]
      # Check that the expected node is still expr_stmt, but now it has 4 children
      # (before comment splicing it had 3), the last child being the comment.
      self._AssertNodeType('expr_stmt', expr)
      self.assertEqual(4, len(expr.children))
      comment_node = expr.children[3]
      self._AssertNodeIsComment(comment_node, '# and a comment')
  
    def testSimpleSeparateLine(self):
      code = textwrap.dedent("""\
          foo = 1
          # first comment
          bar = 2
      """)
      tree = pytree_utils.ParseCodeToTree(code)
      comment_splicer.SpliceComments(tree)
  
      # The comment should've been added to the root's children (now 4, including
      # the ENDMARKER in the end.
      self.assertEqual(4, len(tree.children))
      comment_node = tree.children[1]
      self._AssertNodeIsComment(comment_node)
  
    def testTwoLineComment(self):
      code = textwrap.dedent("""\
          foo = 1
          # first comment
          # second comment
          bar = 2
      """)
      tree = pytree_utils.ParseCodeToTree(code)
      comment_splicer.SpliceComments(tree)
  
      # This is similar to the single-line standalone comment.
      self.assertEqual(4, len(tree.children))
      self._AssertNodeIsComment(tree.children[1])
  
    def testCommentIsFirstChildInCompound(self):
      code = textwrap.dedent("""
          if x:
            # a comment
            foo = 1
      """)
      tree = pytree_utils.ParseCodeToTree(code)
      comment_splicer.SpliceComments(tree)
  
      # Look into the suite node under the 'if'. We don't care about the NEWLINE
      # leaf but the new COMMENT must be a child of the suite and before the
      # actual code leaf.
      if_suite = tree.children[0].children[3]
      self._AssertNodeType('NEWLINE', if_suite.children[0])
      self._AssertNodeIsComment(if_suite.children[1])
  
    def testCommentIsLastChildInCompound(self):
      code = textwrap.dedent("""\
          if x:
            foo = 1
            # a comment
      """)
      tree = pytree_utils.ParseCodeToTree(code)
      comment_splicer.SpliceComments(tree)
  
      # Look into the suite node under the 'if'. We don't care about the DEDENT
      # leaf but the new COMMENT must be a child of the suite and after the
      # actual code leaf.
      if_suite = tree.children[0].children[3]
      self._AssertNodeType('DEDENT', if_suite.children[-1])
      self._AssertNodeIsComment(if_suite.children[-2])
  
    def testInlineAfterSeparateLine(self):
      code = textwrap.dedent("""\
          bar = 1
          # line comment
          foo = 1 # inline comment
      """)
      tree = pytree_utils.ParseCodeToTree(code)
      comment_splicer.SpliceComments(tree)
  
      # The separate line comment should become a child of the root, while
      # the inline comment remains within its simple_node.
      sep_comment_node = tree.children[1]
      self._AssertNodeIsComment(sep_comment_node, '# line comment')
  
      expr = tree.children[2].children[0]
      inline_comment_node = expr.children[-1]
      self._AssertNodeIsComment(inline_comment_node, '# inline comment')
  
    def testSeparateLineAfterInline(self):
      code = textwrap.dedent("""\
          bar = 1
          foo = 1 # inline comment
          # line comment
      """)
      tree = pytree_utils.ParseCodeToTree(code)
      comment_splicer.SpliceComments(tree)
  
      # The separate line comment should become a child of the root, while
      # the inline comment remains within its simple_node.
      sep_comment_node = tree.children[-2]
      self._AssertNodeIsComment(sep_comment_node, '# line comment')
  
      expr = tree.children[1].children[0]
      inline_comment_node = expr.children[-1]
      self._AssertNodeIsComment(inline_comment_node, '# inline comment')
  
    def testCommentBeforeDedent(self):
      code = textwrap.dedent("""\
          if bar:
            z = 1
          # a comment
          j = 2
      """)
      tree = pytree_utils.ParseCodeToTree(code)
      comment_splicer.SpliceComments(tree)
  
      # The comment should go under the tree root, not under the 'if'.
      self._AssertNodeIsComment(tree.children[1])
      if_suite = tree.children[0].children[3]
      self._AssertNodeType('DEDENT', if_suite.children[-1])
  
    def testCommentBeforeDedentTwoLevel(self):
      code = textwrap.dedent("""\
          if foo:
            if bar:
              z = 1
            # a comment
          y = 1
      """)
      tree = pytree_utils.ParseCodeToTree(code)
      comment_splicer.SpliceComments(tree)
  
      if_suite = tree.children[0].children[3]
      # The comment is in the first if_suite, not the nested if under it. It's
      # right before the DEDENT
      self._AssertNodeIsComment(if_suite.children[-2])
      self._AssertNodeType('DEDENT', if_suite.children[-1])
  
    def testCommentBeforeDedentTwoLevelImproperlyIndented(self):
      code = textwrap.dedent("""\
          if foo:
            if bar:
              z = 1
             # comment 2
          y = 1
      """)
      tree = pytree_utils.ParseCodeToTree(code)
      comment_splicer.SpliceComments(tree)
  
      # The comment here is indented by 3 spaces, which is unlike any of the
      # surrounding statement indentation levels. The splicer attaches it to the
      # "closest" parent with smaller indentation.
      if_suite = tree.children[0].children[3]
      # The comment is in the first if_suite, not the nested if under it. It's
      # right before the DEDENT
      self._AssertNodeIsComment(if_suite.children[-2])
      self._AssertNodeType('DEDENT', if_suite.children[-1])
  
    def testCommentBeforeDedentThreeLevel(self):
      code = textwrap.dedent("""\
          if foo:
            if bar:
              z = 1
              # comment 2
            # comment 1
          # comment 0
          j = 2
      """)
      tree = pytree_utils.ParseCodeToTree(code)
      comment_splicer.SpliceComments(tree)
  
      # comment 0 should go under the tree root
      self._AssertNodeIsComment(tree.children[1], '# comment 0')
  
      # comment 1 is in the first if_suite, right before the DEDENT
      if_suite_1 = self._FindNthChildNamed(tree, 'suite', n=1)
      self._AssertNodeIsComment(if_suite_1.children[-2], '# comment 1')
      self._AssertNodeType('DEDENT', if_suite_1.children[-1])
  
      # comment 2 is in if_suite nested under the first if suite,
      # right before the DEDENT
      if_suite_2 = self._FindNthChildNamed(tree, 'suite', n=2)
      self._AssertNodeIsComment(if_suite_2.children[-2], '# comment 2')
      self._AssertNodeType('DEDENT', if_suite_2.children[-1])
  
    def testCommentsInClass(self):
      code = textwrap.dedent("""\
          class Foo:
            '''docstring abc...'''
            # top-level comment
            def foo(): pass
            # another comment
      """)
  
      tree = pytree_utils.ParseCodeToTree(code)
      comment_splicer.SpliceComments(tree)
  
      class_suite = tree.children[0].children[3]
      another_comment = class_suite.children[-2]
      self._AssertNodeIsComment(another_comment, '# another')
  
      # It's OK for the comment to be a child of funcdef, as long as it's
      # the first child and thus comes before the 'def'.
      funcdef = class_suite.children[3]
      toplevel_comment = funcdef.children[0]
      self._AssertNodeIsComment(toplevel_comment, '# top-level')
  
    def testMultipleBlockComments(self):
      code = textwrap.dedent("""\
          # Block comment number 1
  
          # Block comment number 2
          def f():
            pass
      """)
  
      tree = pytree_utils.ParseCodeToTree(code)
      comment_splicer.SpliceComments(tree)
  
      funcdef = tree.children[0]
      block_comment_1 = funcdef.children[0]
      self._AssertNodeIsComment(block_comment_1, '# Block comment number 1')
  
      block_comment_2 = funcdef.children[1]
      self._AssertNodeIsComment(block_comment_2, '# Block comment number 2')
  
    def testCommentsOnDedents(self):
      code = textwrap.dedent("""\
          class Foo(object):
            # A comment for qux.
            def qux(self):
              pass
  
            # Interim comment.
  
            def mux(self):
              pass
      """)
  
      tree = pytree_utils.ParseCodeToTree(code)
      comment_splicer.SpliceComments(tree)
  
      classdef = tree.children[0]
      class_suite = classdef.children[6]
      qux_comment = class_suite.children[1]
      self._AssertNodeIsComment(qux_comment, '# A comment for qux.')
  
      interim_comment = class_suite.children[4]
      self._AssertNodeIsComment(interim_comment, '# Interim comment.')
  
    def testExprComments(self):
      code = textwrap.dedent("""\
          foo( # Request fractions of an hour.
            948.0/3600, 20)
      """)
      tree = pytree_utils.ParseCodeToTree(code)
      comment_splicer.SpliceComments(tree)
  
      trailer = self._FindNthChildNamed(tree, 'trailer', 1)
      comment = trailer.children[1]
      self._AssertNodeIsComment(comment, '# Request fractions of an hour.')
  
    def testMultipleCommentsInOneExpr(self):
      code = textwrap.dedent("""\
          foo( # com 1
            948.0/3600, # com 2
            20 + 12 # com 3
            )
      """)
      tree = pytree_utils.ParseCodeToTree(code)
      comment_splicer.SpliceComments(tree)
  
      trailer = self._FindNthChildNamed(tree, 'trailer', 1)
      self._AssertNodeIsComment(trailer.children[1], '# com 1')
  
      arglist = self._FindNthChildNamed(tree, 'arglist', 1)
      self._AssertNodeIsComment(arglist.children[2], '# com 2')
  
      arith_expr = self._FindNthChildNamed(tree, 'arith_expr', 1)
      self._AssertNodeIsComment(arith_expr.children[-1], '# com 3')
  
  
  if __name__ == '__main__':
    unittest.main()
  