1 # Copyright 2015 Google Inc. All Rights Reserved. 2 # 3 # Licensed under the Apache License, Version 2.0 (the "License"); 4 # you may not use this file except in compliance with the License. 5 # You may obtain a copy of the License at 6 # 7 # http://www.apache.org/licenses/LICENSE-2.0 8 # 9 # Unless required by applicable law or agreed to in writing, software 10 # distributed under the License is distributed on an "AS IS" BASIS, 11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 # See the License for the specific language governing permissions and 13 # limitations under the License. 14 """Tests for yapf.comment_splicer.""" 15 16 import textwrap 17 import unittest 18 19 from yapf.yapflib import comment_splicer 20 from yapf.yapflib import py3compat 21 from yapf.yapflib import pytree_utils 22 23 24 class CommentSplicerTest(unittest.TestCase): 25 26 def _AssertNodeType(self, expected_type, node): 27 self.assertEqual(expected_type, pytree_utils.NodeName(node)) 28 29 def _AssertNodeIsComment(self, node, text_in_comment=None): 30 if pytree_utils.NodeName(node) == 'simple_stmt': 31 self._AssertNodeType('COMMENT', node.children[0]) 32 node_value = node.children[0].value 33 else: 34 self._AssertNodeType('COMMENT', node) 35 node_value = node.value 36 if text_in_comment is not None: 37 self.assertIn(text_in_comment, node_value) 38 39 def _FindNthChildNamed(self, node, name, n=1): 40 for i, child in enumerate( 41 py3compat.ifilter(lambda c: pytree_utils.NodeName(c) == name, 42 node.pre_order())): 43 if i == n - 1: 44 return child 45 raise RuntimeError('No Nth child for n={0}'.format(n)) 46 47 def testSimpleInline(self): 48 code = 'foo = 1 # and a comment\n' 49 tree = pytree_utils.ParseCodeToTree(code) 50 comment_splicer.SpliceComments(tree) 51 52 expr = tree.children[0].children[0] 53 # Check that the expected node is still expr_stmt, but now it has 4 children 54 # (before comment splicing it had 3), the last child being the comment. 55 self._AssertNodeType('expr_stmt', expr) 56 self.assertEqual(4, len(expr.children)) 57 comment_node = expr.children[3] 58 self._AssertNodeIsComment(comment_node, '# and a comment') 59 60 def testSimpleSeparateLine(self): 61 code = textwrap.dedent(r''' 62 foo = 1 63 # first comment 64 bar = 2 65 ''') 66 tree = pytree_utils.ParseCodeToTree(code) 67 comment_splicer.SpliceComments(tree) 68 69 # The comment should've been added to the root's children (now 4, including 70 # the ENDMARKER in the end. 71 self.assertEqual(4, len(tree.children)) 72 comment_node = tree.children[1] 73 self._AssertNodeIsComment(comment_node) 74 75 def testTwoLineComment(self): 76 code = textwrap.dedent(r''' 77 foo = 1 78 # first comment 79 # second comment 80 bar = 2 81 ''') 82 tree = pytree_utils.ParseCodeToTree(code) 83 comment_splicer.SpliceComments(tree) 84 85 # This is similar to the single-line standalone comment. 86 self.assertEqual(4, len(tree.children)) 87 self._AssertNodeIsComment(tree.children[1]) 88 89 def testCommentIsFirstChildInCompound(self): 90 code = textwrap.dedent(r''' 91 if x: 92 # a comment 93 foo = 1 94 ''') 95 tree = pytree_utils.ParseCodeToTree(code) 96 comment_splicer.SpliceComments(tree) 97 98 # Look into the suite node under the 'if'. We don't care about the NEWLINE 99 # leaf but the new COMMENT must be a child of the suite and before the 100 # actual code leaf. 101 if_suite = tree.children[0].children[3] 102 self._AssertNodeType('NEWLINE', if_suite.children[0]) 103 self._AssertNodeIsComment(if_suite.children[1]) 104 105 def testCommentIsLastChildInCompound(self): 106 code = textwrap.dedent(r''' 107 if x: 108 foo = 1 109 # a comment 110 ''') 111 tree = pytree_utils.ParseCodeToTree(code) 112 comment_splicer.SpliceComments(tree) 113 114 # Look into the suite node under the 'if'. We don't care about the DEDENT 115 # leaf but the new COMMENT must be a child of the suite and after the 116 # actual code leaf. 117 if_suite = tree.children[0].children[3] 118 self._AssertNodeType('DEDENT', if_suite.children[-1]) 119 self._AssertNodeIsComment(if_suite.children[-2]) 120 121 def testInlineAfterSeparateLine(self): 122 code = textwrap.dedent(r''' 123 bar = 1 124 # line comment 125 foo = 1 # inline comment 126 ''') 127 tree = pytree_utils.ParseCodeToTree(code) 128 comment_splicer.SpliceComments(tree) 129 130 # The separate line comment should become a child of the root, while 131 # the inline comment remains within its simple_node. 132 sep_comment_node = tree.children[1] 133 self._AssertNodeIsComment(sep_comment_node, '# line comment') 134 135 expr = tree.children[2].children[0] 136 inline_comment_node = expr.children[-1] 137 self._AssertNodeIsComment(inline_comment_node, '# inline comment') 138 139 def testSeparateLineAfterInline(self): 140 code = textwrap.dedent(r''' 141 bar = 1 142 foo = 1 # inline comment 143 # line comment 144 ''') 145 tree = pytree_utils.ParseCodeToTree(code) 146 comment_splicer.SpliceComments(tree) 147 148 # The separate line comment should become a child of the root, while 149 # the inline comment remains within its simple_node. 150 sep_comment_node = tree.children[-2] 151 self._AssertNodeIsComment(sep_comment_node, '# line comment') 152 153 expr = tree.children[1].children[0] 154 inline_comment_node = expr.children[-1] 155 self._AssertNodeIsComment(inline_comment_node, '# inline comment') 156 157 def testCommentBeforeDedent(self): 158 code = textwrap.dedent(r''' 159 if bar: 160 z = 1 161 # a comment 162 j = 2 163 ''') 164 tree = pytree_utils.ParseCodeToTree(code) 165 comment_splicer.SpliceComments(tree) 166 167 # The comment should go under the tree root, not under the 'if'. 168 self._AssertNodeIsComment(tree.children[1]) 169 if_suite = tree.children[0].children[3] 170 self._AssertNodeType('DEDENT', if_suite.children[-1]) 171 172 def testCommentBeforeDedentTwoLevel(self): 173 code = textwrap.dedent(r''' 174 if foo: 175 if bar: 176 z = 1 177 # a comment 178 y = 1 179 ''') 180 tree = pytree_utils.ParseCodeToTree(code) 181 comment_splicer.SpliceComments(tree) 182 183 if_suite = tree.children[0].children[3] 184 # The comment is in the first if_suite, not the nested if under it. It's 185 # right before the DEDENT 186 self._AssertNodeIsComment(if_suite.children[-2]) 187 self._AssertNodeType('DEDENT', if_suite.children[-1]) 188 189 def testCommentBeforeDedentTwoLevelImproperlyIndented(self): 190 code = textwrap.dedent(r''' 191 if foo: 192 if bar: 193 z = 1 194 # comment 2 195 y = 1 196 ''') 197 tree = pytree_utils.ParseCodeToTree(code) 198 comment_splicer.SpliceComments(tree) 199 200 # The comment here is indented by 3 spaces, which is unlike any of the 201 # surrounding statement indentation levels. The splicer attaches it to the 202 # "closest" parent with smaller indentation. 203 if_suite = tree.children[0].children[3] 204 # The comment is in the first if_suite, not the nested if under it. It's 205 # right before the DEDENT 206 self._AssertNodeIsComment(if_suite.children[-2]) 207 self._AssertNodeType('DEDENT', if_suite.children[-1]) 208 209 def testCommentBeforeDedentThreeLevel(self): 210 code = textwrap.dedent(r''' 211 if foo: 212 if bar: 213 z = 1 214 # comment 2 215 # comment 1 216 # comment 0 217 j = 2 218 ''') 219 tree = pytree_utils.ParseCodeToTree(code) 220 comment_splicer.SpliceComments(tree) 221 222 # comment 0 should go under the tree root 223 self._AssertNodeIsComment(tree.children[1], '# comment 0') 224 225 # comment 1 is in the first if_suite, right before the DEDENT 226 if_suite_1 = self._FindNthChildNamed(tree, 'suite', n=1) 227 self._AssertNodeIsComment(if_suite_1.children[-2], '# comment 1') 228 self._AssertNodeType('DEDENT', if_suite_1.children[-1]) 229 230 # comment 2 is in if_suite nested under the first if suite, 231 # right before the DEDENT 232 if_suite_2 = self._FindNthChildNamed(tree, 'suite', n=2) 233 self._AssertNodeIsComment(if_suite_2.children[-2], '# comment 2') 234 self._AssertNodeType('DEDENT', if_suite_2.children[-1]) 235 236 def testCommentsInClass(self): 237 code = textwrap.dedent(r''' 238 class Foo: 239 """docstring abc...""" 240 # top-level comment 241 def foo(): pass 242 # another comment 243 ''') 244 245 tree = pytree_utils.ParseCodeToTree(code) 246 comment_splicer.SpliceComments(tree) 247 248 class_suite = tree.children[0].children[3] 249 another_comment = class_suite.children[-2] 250 self._AssertNodeIsComment(another_comment, '# another') 251 252 # It's OK for the comment to be a child of funcdef, as long as it's 253 # the first child and thus comes before the 'def'. 254 funcdef = class_suite.children[3] 255 toplevel_comment = funcdef.children[0] 256 self._AssertNodeIsComment(toplevel_comment, '# top-level') 257 258 def testMultipleBlockComments(self): 259 code = textwrap.dedent(r''' 260 # Block comment number 1 261 262 # Block comment number 2 263 def f(): 264 pass 265 ''') 266 267 tree = pytree_utils.ParseCodeToTree(code) 268 comment_splicer.SpliceComments(tree) 269 270 funcdef = tree.children[0] 271 block_comment_1 = funcdef.children[0] 272 self._AssertNodeIsComment(block_comment_1, '# Block comment number 1') 273 274 block_comment_2 = funcdef.children[1] 275 self._AssertNodeIsComment(block_comment_2, '# Block comment number 2') 276 277 def testCommentsOnDedents(self): 278 code = textwrap.dedent(r''' 279 class Foo(object): 280 # A comment for qux. 281 def qux(self): 282 pass 283 284 # Interim comment. 285 286 def mux(self): 287 pass 288 ''') 289 290 tree = pytree_utils.ParseCodeToTree(code) 291 comment_splicer.SpliceComments(tree) 292 293 classdef = tree.children[0] 294 class_suite = classdef.children[6] 295 qux_comment = class_suite.children[1] 296 self._AssertNodeIsComment(qux_comment, '# A comment for qux.') 297 298 interim_comment = class_suite.children[4] 299 self._AssertNodeIsComment(interim_comment, '# Interim comment.') 300 301 def testExprComments(self): 302 code = textwrap.dedent(r''' 303 foo( # Request fractions of an hour. 304 948.0/3600, 20) 305 ''') 306 tree = pytree_utils.ParseCodeToTree(code) 307 comment_splicer.SpliceComments(tree) 308 309 trailer = self._FindNthChildNamed(tree, 'trailer', 1) 310 comment = trailer.children[1] 311 self._AssertNodeIsComment(comment, '# Request fractions of an hour.') 312 313 def testMultipleCommentsInOneExpr(self): 314 code = textwrap.dedent(r''' 315 foo( # com 1 316 948.0/3600, # com 2 317 20 + 12 # com 3 318 ) 319 ''') 320 tree = pytree_utils.ParseCodeToTree(code) 321 comment_splicer.SpliceComments(tree) 322 323 trailer = self._FindNthChildNamed(tree, 'trailer', 1) 324 self._AssertNodeIsComment(trailer.children[1], '# com 1') 325 326 arglist = self._FindNthChildNamed(tree, 'arglist', 1) 327 self._AssertNodeIsComment(arglist.children[2], '# com 2') 328 329 arith_expr = self._FindNthChildNamed(tree, 'arith_expr', 1) 330 self._AssertNodeIsComment(arith_expr.children[-1], '# com 3') 331 332 333 if __name__ == '__main__': 334 unittest.main() 335
