| # 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 |
| |
| |
| class CommentSplicerTest(unittest.TestCase): |
| |
| 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 = 'foo = 1 # and a comment\n' |
| 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(r''' |
| 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(r''' |
| 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(r''' |
| 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(r''' |
| 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(r''' |
| 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(r''' |
| 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(r''' |
| 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(r''' |
| 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(r''' |
| 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(r''' |
| 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(r''' |
| 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(r''' |
| # 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(r''' |
| 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(r''' |
| 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(r''' |
| 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() |