blob: b297560d86c8cab8b9c9f611d7af9fea5f82a021 [file] [log] [blame]
Bill Wendling92290f12018-01-14 16:49:35 -08001# Copyright 2015 Google Inc. All Rights Reserved.
Bill Wendling7d623452015-03-18 13:36:07 -07002#
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.pytree_unwrapper."""
15
Bill Wendling7d623452015-03-18 13:36:07 -070016import textwrap
17import unittest
18
Bill Wendling831fb3d2021-12-27 13:04:22 -080019from yapf.pytree import pytree_utils
Bill Wendling7d623452015-03-18 13:36:07 -070020
Bill Wendlinga8ebae22016-10-23 16:09:06 -070021from yapftests import yapf_test_helper
Bill Wendling7d623452015-03-18 13:36:07 -070022
Bill Wendling7d623452015-03-18 13:36:07 -070023
Bill Wendlinga8ebae22016-10-23 16:09:06 -070024class PytreeUnwrapperTest(yapf_test_helper.YAPFTest):
Bill Wendling7d623452015-03-18 13:36:07 -070025
Bill Wendling5b75e6f2021-12-25 22:31:14 -060026 def _CheckLogicalLines(self, llines, list_of_expected):
27 """Check that the given LogicalLines match expectations.
Bill Wendling7d623452015-03-18 13:36:07 -070028
29 Args:
Bill Wendling5b75e6f2021-12-25 22:31:14 -060030 llines: list of LogicalLine
Bill Wendling7d623452015-03-18 13:36:07 -070031 list_of_expected: list of (depth, values) pairs. Non-semantic tokens are
32 filtered out from the expected values.
33 """
34 actual = []
Bill Wendling5b75e6f2021-12-25 22:31:14 -060035 for lline in llines:
Bill Wendlingeef8fbd2016-10-05 00:34:25 -070036 filtered_values = [
Matthew Suozzoba8f51f2017-10-31 00:33:38 -040037 ft.value
Bill Wendling5b75e6f2021-12-25 22:31:14 -060038 for ft in lline.tokens
Bill Wendlingeef8fbd2016-10-05 00:34:25 -070039 if ft.name not in pytree_utils.NONSEMANTIC_TOKENS
40 ]
Bill Wendling5b75e6f2021-12-25 22:31:14 -060041 actual.append((lline.depth, filtered_values))
Bill Wendling7d623452015-03-18 13:36:07 -070042
43 self.assertEqual(list_of_expected, actual)
44
45 def testSimpleFileScope(self):
Bill Wendlinga0db84c2023-06-23 13:24:34 -050046 code = textwrap.dedent("""\
47 x = 1
48 # a comment
49 y = 2
50 """)
Bill Wendling5b75e6f2021-12-25 22:31:14 -060051 llines = yapf_test_helper.ParseAndUnwrap(code)
52 self._CheckLogicalLines(llines, [
Bill Wendling7d623452015-03-18 13:36:07 -070053 (0, ['x', '=', '1']),
54 (0, ['# a comment']),
Bill Wendlinga2f03422016-10-23 16:36:47 -070055 (0, ['y', '=', '2']),
56 ])
Bill Wendling7d623452015-03-18 13:36:07 -070057
58 def testSimpleMultilineStatement(self):
Bill Wendlinga0db84c2023-06-23 13:24:34 -050059 code = textwrap.dedent("""\
60 y = (1 +
61 x)
62 """)
Bill Wendling5b75e6f2021-12-25 22:31:14 -060063 llines = yapf_test_helper.ParseAndUnwrap(code)
64 self._CheckLogicalLines(llines, [
Bill Wendlinga2f03422016-10-23 16:36:47 -070065 (0, ['y', '=', '(', '1', '+', 'x', ')']),
66 ])
Bill Wendling7d623452015-03-18 13:36:07 -070067
68 def testFileScopeWithInlineComment(self):
Bill Wendlinga0db84c2023-06-23 13:24:34 -050069 code = textwrap.dedent("""\
70 x = 1 # a comment
71 y = 2
72 """)
Bill Wendling5b75e6f2021-12-25 22:31:14 -060073 llines = yapf_test_helper.ParseAndUnwrap(code)
74 self._CheckLogicalLines(llines, [
Bill Wendling7d623452015-03-18 13:36:07 -070075 (0, ['x', '=', '1', '# a comment']),
Bill Wendlinga2f03422016-10-23 16:36:47 -070076 (0, ['y', '=', '2']),
77 ])
Bill Wendling7d623452015-03-18 13:36:07 -070078
79 def testSimpleIf(self):
Bill Wendlinga0db84c2023-06-23 13:24:34 -050080 code = textwrap.dedent("""\
81 if foo:
82 x = 1
83 y = 2
84 """)
Bill Wendling5b75e6f2021-12-25 22:31:14 -060085 llines = yapf_test_helper.ParseAndUnwrap(code)
86 self._CheckLogicalLines(llines, [
Bill Wendling7d623452015-03-18 13:36:07 -070087 (0, ['if', 'foo', ':']),
88 (1, ['x', '=', '1']),
Bill Wendlinga2f03422016-10-23 16:36:47 -070089 (1, ['y', '=', '2']),
90 ])
Bill Wendling7d623452015-03-18 13:36:07 -070091
92 def testSimpleIfWithComments(self):
Bill Wendlinga0db84c2023-06-23 13:24:34 -050093 code = textwrap.dedent("""\
94 # c1
95 if foo: # c2
96 x = 1
97 y = 2
98 """)
Bill Wendling5b75e6f2021-12-25 22:31:14 -060099 llines = yapf_test_helper.ParseAndUnwrap(code)
100 self._CheckLogicalLines(llines, [
Bill Wendling7d623452015-03-18 13:36:07 -0700101 (0, ['# c1']),
102 (0, ['if', 'foo', ':', '# c2']),
103 (1, ['x', '=', '1']),
Bill Wendlinga2f03422016-10-23 16:36:47 -0700104 (1, ['y', '=', '2']),
105 ])
Bill Wendling7d623452015-03-18 13:36:07 -0700106
107 def testIfWithCommentsInside(self):
Bill Wendlinga0db84c2023-06-23 13:24:34 -0500108 code = textwrap.dedent("""\
109 if foo:
110 # c1
111 x = 1 # c2
112 # c3
113 y = 2
114 """)
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600115 llines = yapf_test_helper.ParseAndUnwrap(code)
116 self._CheckLogicalLines(llines, [
Bill Wendling7d623452015-03-18 13:36:07 -0700117 (0, ['if', 'foo', ':']),
118 (1, ['# c1']),
119 (1, ['x', '=', '1', '# c2']),
120 (1, ['# c3']),
Bill Wendlinga2f03422016-10-23 16:36:47 -0700121 (1, ['y', '=', '2']),
122 ])
Bill Wendling7d623452015-03-18 13:36:07 -0700123
124 def testIfElifElse(self):
Bill Wendlinga0db84c2023-06-23 13:24:34 -0500125 code = textwrap.dedent("""\
126 if x:
127 x = 1 # c1
128 elif y: # c2
129 y = 1
130 else:
131 # c3
132 z = 1
133 """)
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600134 llines = yapf_test_helper.ParseAndUnwrap(code)
135 self._CheckLogicalLines(llines, [
Bill Wendling7d623452015-03-18 13:36:07 -0700136 (0, ['if', 'x', ':']),
137 (1, ['x', '=', '1', '# c1']),
138 (0, ['elif', 'y', ':', '# c2']),
139 (1, ['y', '=', '1']),
140 (0, ['else', ':']),
141 (1, ['# c3']),
Bill Wendlinga2f03422016-10-23 16:36:47 -0700142 (1, ['z', '=', '1']),
143 ])
Bill Wendling7d623452015-03-18 13:36:07 -0700144
145 def testNestedCompoundTwoLevel(self):
Bill Wendlinga0db84c2023-06-23 13:24:34 -0500146 code = textwrap.dedent("""\
147 if x:
148 x = 1 # c1
149 while t:
150 # c2
151 j = 1
152 k = 1
153 """)
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600154 llines = yapf_test_helper.ParseAndUnwrap(code)
155 self._CheckLogicalLines(llines, [
Bill Wendling7d623452015-03-18 13:36:07 -0700156 (0, ['if', 'x', ':']),
157 (1, ['x', '=', '1', '# c1']),
158 (1, ['while', 't', ':']),
159 (2, ['# c2']),
160 (2, ['j', '=', '1']),
Bill Wendlinga2f03422016-10-23 16:36:47 -0700161 (1, ['k', '=', '1']),
162 ])
Bill Wendling7d623452015-03-18 13:36:07 -0700163
164 def testSimpleWhile(self):
Bill Wendlinga0db84c2023-06-23 13:24:34 -0500165 code = textwrap.dedent("""\
166 while x > 1: # c1
167 # c2
168 x = 1
169 """)
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600170 llines = yapf_test_helper.ParseAndUnwrap(code)
171 self._CheckLogicalLines(llines, [
Bill Wendling7d623452015-03-18 13:36:07 -0700172 (0, ['while', 'x', '>', '1', ':', '# c1']),
173 (1, ['# c2']),
Bill Wendlinga2f03422016-10-23 16:36:47 -0700174 (1, ['x', '=', '1']),
175 ])
Bill Wendling7d623452015-03-18 13:36:07 -0700176
177 def testSimpleTry(self):
Bill Wendlinga0db84c2023-06-23 13:24:34 -0500178 code = textwrap.dedent("""\
179 try:
180 pass
181 except:
182 pass
183 except:
184 pass
185 else:
186 pass
187 finally:
188 pass
189 """)
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600190 llines = yapf_test_helper.ParseAndUnwrap(code)
191 self._CheckLogicalLines(llines, [
Bill Wendling7d623452015-03-18 13:36:07 -0700192 (0, ['try', ':']),
193 (1, ['pass']),
194 (0, ['except', ':']),
195 (1, ['pass']),
196 (0, ['except', ':']),
197 (1, ['pass']),
198 (0, ['else', ':']),
199 (1, ['pass']),
200 (0, ['finally', ':']),
Bill Wendlinga2f03422016-10-23 16:36:47 -0700201 (1, ['pass']),
202 ])
Bill Wendling7d623452015-03-18 13:36:07 -0700203
204 def testSimpleFuncdef(self):
Bill Wendlinga0db84c2023-06-23 13:24:34 -0500205 code = textwrap.dedent("""\
206 def foo(x): # c1
207 # c2
208 return x
209 """)
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600210 llines = yapf_test_helper.ParseAndUnwrap(code)
211 self._CheckLogicalLines(llines, [
Bill Wendling7d623452015-03-18 13:36:07 -0700212 (0, ['def', 'foo', '(', 'x', ')', ':', '# c1']),
213 (1, ['# c2']),
Bill Wendlinga2f03422016-10-23 16:36:47 -0700214 (1, ['return', 'x']),
215 ])
Bill Wendling7d623452015-03-18 13:36:07 -0700216
217 def testTwoFuncDefs(self):
Bill Wendlinga0db84c2023-06-23 13:24:34 -0500218 code = textwrap.dedent("""\
219 def foo(x): # c1
220 # c2
221 return x
Bill Wendling7d623452015-03-18 13:36:07 -0700222
Bill Wendlinga0db84c2023-06-23 13:24:34 -0500223 def bar(): # c3
224 # c4
225 return x
226 """)
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600227 llines = yapf_test_helper.ParseAndUnwrap(code)
228 self._CheckLogicalLines(llines, [
Bill Wendling7d623452015-03-18 13:36:07 -0700229 (0, ['def', 'foo', '(', 'x', ')', ':', '# c1']),
230 (1, ['# c2']),
231 (1, ['return', 'x']),
232 (0, ['def', 'bar', '(', ')', ':', '# c3']),
233 (1, ['# c4']),
Bill Wendlinga2f03422016-10-23 16:36:47 -0700234 (1, ['return', 'x']),
235 ])
Bill Wendling7d623452015-03-18 13:36:07 -0700236
237 def testSimpleClassDef(self):
Bill Wendlinga0db84c2023-06-23 13:24:34 -0500238 code = textwrap.dedent("""\
239 class Klass: # c1
240 # c2
241 p = 1
242 """)
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600243 llines = yapf_test_helper.ParseAndUnwrap(code)
244 self._CheckLogicalLines(llines, [
Bill Wendling7d623452015-03-18 13:36:07 -0700245 (0, ['class', 'Klass', ':', '# c1']),
246 (1, ['# c2']),
Bill Wendlinga2f03422016-10-23 16:36:47 -0700247 (1, ['p', '=', '1']),
248 ])
Bill Wendling7d623452015-03-18 13:36:07 -0700249
250 def testSingleLineStmtInFunc(self):
Bill Wendlinga0db84c2023-06-23 13:24:34 -0500251 code = textwrap.dedent("""\
Bill Wendling7d623452015-03-18 13:36:07 -0700252 def f(): return 37
Bill Wendlinga0db84c2023-06-23 13:24:34 -0500253 """)
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600254 llines = yapf_test_helper.ParseAndUnwrap(code)
255 self._CheckLogicalLines(llines, [
Bill Wendling7d623452015-03-18 13:36:07 -0700256 (0, ['def', 'f', '(', ')', ':']),
Bill Wendlinga2f03422016-10-23 16:36:47 -0700257 (1, ['return', '37']),
258 ])
Bill Wendling7d623452015-03-18 13:36:07 -0700259
260 def testMultipleComments(self):
Bill Wendlinga0db84c2023-06-23 13:24:34 -0500261 code = textwrap.dedent("""\
Bill Wendling7d623452015-03-18 13:36:07 -0700262 # Comment #1
263
264 # Comment #2
265 def f():
266 pass
Bill Wendlinga0db84c2023-06-23 13:24:34 -0500267 """)
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600268 llines = yapf_test_helper.ParseAndUnwrap(code)
269 self._CheckLogicalLines(llines, [
Bill Wendling7d623452015-03-18 13:36:07 -0700270 (0, ['# Comment #1']),
271 (0, ['# Comment #2']),
272 (0, ['def', 'f', '(', ')', ':']),
Bill Wendlinga2f03422016-10-23 16:36:47 -0700273 (1, ['pass']),
274 ])
Bill Wendling7d623452015-03-18 13:36:07 -0700275
Bill Wendling5b419822015-07-06 01:39:50 -0700276 def testSplitListWithComment(self):
Bill Wendlinga0db84c2023-06-23 13:24:34 -0500277 code = textwrap.dedent("""\
278 a = [
279 'a',
280 'b',
281 'c', # hello world
282 ]
283 """)
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600284 llines = yapf_test_helper.ParseAndUnwrap(code)
285 self._CheckLogicalLines(llines, [(0, [
Bill Wendlinga2f03422016-10-23 16:36:47 -0700286 'a', '=', '[', "'a'", ',', "'b'", ',', "'c'", ',', '# hello world', ']'
287 ])])
Bill Wendling5b419822015-07-06 01:39:50 -0700288
Bill Wendling7d623452015-03-18 13:36:07 -0700289
Bill Wendlinga8ebae22016-10-23 16:09:06 -0700290class MatchBracketsTest(yapf_test_helper.YAPFTest):
Bill Wendling7d623452015-03-18 13:36:07 -0700291
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600292 def _CheckMatchingBrackets(self, llines, list_of_expected):
Bill Wendling7d623452015-03-18 13:36:07 -0700293 """Check that the tokens have the expected matching bracket.
294
295 Arguments:
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600296 llines: list of LogicalLine.
Bill Wendling7d623452015-03-18 13:36:07 -0700297 list_of_expected: list of (index, index) pairs. The matching brackets at
298 the indexes need to match. Non-semantic tokens are filtered out from the
299 expected values.
300 """
301 actual = []
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600302 for lline in llines:
Matthew Suozzoba8f51f2017-10-31 00:33:38 -0400303 filtered_values = [(ft, ft.matching_bracket)
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600304 for ft in lline.tokens
Bill Wendling7d623452015-03-18 13:36:07 -0700305 if ft.name not in pytree_utils.NONSEMANTIC_TOKENS]
306 if filtered_values:
307 actual.append(filtered_values)
308
309 for index, bracket_list in enumerate(list_of_expected):
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600310 lline = actual[index]
Bill Wendling7d623452015-03-18 13:36:07 -0700311 if not bracket_list:
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600312 for value in lline:
Bill Wendling7d623452015-03-18 13:36:07 -0700313 self.assertIsNone(value[1])
314 else:
315 for open_bracket, close_bracket in bracket_list:
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600316 self.assertEqual(lline[open_bracket][0], lline[close_bracket][1])
317 self.assertEqual(lline[close_bracket][0], lline[open_bracket][1])
Bill Wendling7d623452015-03-18 13:36:07 -0700318
319 def testFunctionDef(self):
320 code = textwrap.dedent("""\
Bill Wendling2b0c6312018-01-13 21:53:14 -0800321 def foo(a, b=['w','d'], c=[42, 37]):
Bill Wendling7d623452015-03-18 13:36:07 -0700322 pass
Bill Wendlinga0db84c2023-06-23 13:24:34 -0500323 """)
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600324 llines = yapf_test_helper.ParseAndUnwrap(code)
325 self._CheckMatchingBrackets(llines, [
Bill Wendling2b0c6312018-01-13 21:53:14 -0800326 [(2, 20), (7, 11), (15, 19)],
Bill Wendlinga2f03422016-10-23 16:36:47 -0700327 [],
328 ])
Bill Wendling7d623452015-03-18 13:36:07 -0700329
330 def testDecorator(self):
331 code = textwrap.dedent("""\
332 @bar()
333 def foo(a, b, c):
334 pass
Bill Wendlinga0db84c2023-06-23 13:24:34 -0500335 """)
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600336 llines = yapf_test_helper.ParseAndUnwrap(code)
337 self._CheckMatchingBrackets(llines, [
Bill Wendling7d623452015-03-18 13:36:07 -0700338 [(2, 3)],
339 [(2, 8)],
Bill Wendlinga2f03422016-10-23 16:36:47 -0700340 [],
341 ])
Bill Wendling7d623452015-03-18 13:36:07 -0700342
343 def testClassDef(self):
344 code = textwrap.dedent("""\
345 class A(B, C, D):
346 pass
Bill Wendlinga0db84c2023-06-23 13:24:34 -0500347 """)
Bill Wendling5b75e6f2021-12-25 22:31:14 -0600348 llines = yapf_test_helper.ParseAndUnwrap(code)
349 self._CheckMatchingBrackets(llines, [
Bill Wendling7d623452015-03-18 13:36:07 -0700350 [(2, 8)],
Bill Wendlinga2f03422016-10-23 16:36:47 -0700351 [],
352 ])
Bill Wendling7d623452015-03-18 13:36:07 -0700353
354
Bill Wendling7d623452015-03-18 13:36:07 -0700355if __name__ == '__main__':
356 unittest.main()