blob: 64cb7083e61524fa699969ade0aaa4676541b4a4 [file] [log] [blame]
# 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()