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