• 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.yapflib import comment_splicer
20from yapf.yapflib import py3compat
21from yapf.yapflib import pytree_utils
22
23
24class 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
333if __name__ == '__main__':
334  unittest.main()
335