• 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.split_penalty."""
15
16import sys
17import textwrap
18import unittest
19
20from lib2to3 import pytree
21
22from yapf.yapflib import pytree_utils
23from yapf.yapflib import pytree_visitor
24from yapf.yapflib import split_penalty
25
26UNBREAKABLE = split_penalty.UNBREAKABLE
27VERY_STRONGLY_CONNECTED = split_penalty.VERY_STRONGLY_CONNECTED
28DOTTED_NAME = split_penalty.DOTTED_NAME
29STRONGLY_CONNECTED = split_penalty.STRONGLY_CONNECTED
30
31
32class SplitPenaltyTest(unittest.TestCase):
33
34  def _ParseAndComputePenalties(self, code, dumptree=False):
35    """Parses the code and computes split penalties.
36
37    Arguments:
38      code: code to parse as a string
39      dumptree: if True, the parsed pytree (after penalty assignment) is dumped
40        to stderr. Useful for debugging.
41
42    Returns:
43      Parse tree.
44    """
45    tree = pytree_utils.ParseCodeToTree(code)
46    split_penalty.ComputeSplitPenalties(tree)
47    if dumptree:
48      pytree_visitor.DumpPyTree(tree, target_stream=sys.stderr)
49    return tree
50
51  def _CheckPenalties(self, tree, list_of_expected):
52    """Check that the tokens in the tree have the correct penalties.
53
54    Args:
55      tree: the pytree.
56      list_of_expected: list of (name, penalty) pairs. Non-semantic tokens are
57        filtered out from the expected values.
58    """
59
60    def FlattenRec(tree):
61      if pytree_utils.NodeName(tree) in pytree_utils.NONSEMANTIC_TOKENS:
62        return []
63      if isinstance(tree, pytree.Leaf):
64        return [(tree.value,
65                 pytree_utils.GetNodeAnnotation(
66                     tree, pytree_utils.Annotation.SPLIT_PENALTY))]
67      nodes = []
68      for node in tree.children:
69        nodes += FlattenRec(node)
70      return nodes
71
72    self.assertEqual(list_of_expected, FlattenRec(tree))
73
74  def testUnbreakable(self):
75    # Test function definitions.
76    code = textwrap.dedent(r"""
77      def foo(x):
78        pass
79      """)
80    tree = self._ParseAndComputePenalties(code)
81    self._CheckPenalties(tree, [
82        ('def', None),
83        ('foo', UNBREAKABLE),
84        ('(', UNBREAKABLE),
85        ('x', None),
86        (')', STRONGLY_CONNECTED),
87        (':', UNBREAKABLE),
88        ('pass', None),
89    ])
90
91    # Test function definition with trailing comment.
92    code = textwrap.dedent(r"""
93      def foo(x):  # trailing comment
94        pass
95      """)
96    tree = self._ParseAndComputePenalties(code)
97    self._CheckPenalties(tree, [
98        ('def', None),
99        ('foo', UNBREAKABLE),
100        ('(', UNBREAKABLE),
101        ('x', None),
102        (')', STRONGLY_CONNECTED),
103        (':', UNBREAKABLE),
104        ('pass', None),
105    ])
106
107    # Test class definitions.
108    code = textwrap.dedent(r"""
109      class A:
110        pass
111      class B(A):
112        pass
113      """)
114    tree = self._ParseAndComputePenalties(code)
115    self._CheckPenalties(tree, [
116        ('class', None),
117        ('A', UNBREAKABLE),
118        (':', UNBREAKABLE),
119        ('pass', None),
120        ('class', None),
121        ('B', UNBREAKABLE),
122        ('(', UNBREAKABLE),
123        ('A', None),
124        (')', None),
125        (':', UNBREAKABLE),
126        ('pass', None),
127    ])
128
129    # Test lambda definitions.
130    code = textwrap.dedent(r"""
131      lambda a, b: None
132      """)
133    tree = self._ParseAndComputePenalties(code)
134    self._CheckPenalties(tree, [
135        ('lambda', None),
136        ('a', UNBREAKABLE),
137        (',', UNBREAKABLE),
138        ('b', UNBREAKABLE),
139        (':', UNBREAKABLE),
140        ('None', UNBREAKABLE),
141    ])
142
143    # Test dotted names.
144    code = textwrap.dedent(r"""
145      import a.b.c
146      """)
147    tree = self._ParseAndComputePenalties(code)
148    self._CheckPenalties(tree, [
149        ('import', None),
150        ('a', None),
151        ('.', UNBREAKABLE),
152        ('b', UNBREAKABLE),
153        ('.', UNBREAKABLE),
154        ('c', UNBREAKABLE),
155    ])
156
157  def testStronglyConnected(self):
158    # Test dictionary keys.
159    code = textwrap.dedent(r"""
160      a = {
161          'x': 42,
162          y(lambda a: 23): 37,
163      }
164      """)
165    tree = self._ParseAndComputePenalties(code)
166    self._CheckPenalties(tree, [
167        ('a', None),
168        ('=', None),
169        ('{', None),
170        ("'x'", None),
171        (':', STRONGLY_CONNECTED),
172        ('42', None),
173        (',', None),
174        ('y', None),
175        ('(', UNBREAKABLE),
176        ('lambda', STRONGLY_CONNECTED),
177        ('a', UNBREAKABLE),
178        (':', UNBREAKABLE),
179        ('23', UNBREAKABLE),
180        (')', VERY_STRONGLY_CONNECTED),
181        (':', STRONGLY_CONNECTED),
182        ('37', None),
183        (',', None),
184        ('}', None),
185    ])
186
187    # Test list comprehension.
188    code = textwrap.dedent(r"""
189      [a for a in foo if a.x == 37]
190      """)
191    tree = self._ParseAndComputePenalties(code)
192    self._CheckPenalties(tree, [
193        ('[', None),
194        ('a', None),
195        ('for', 0),
196        ('a', STRONGLY_CONNECTED),
197        ('in', STRONGLY_CONNECTED),
198        ('foo', STRONGLY_CONNECTED),
199        ('if', 0),
200        ('a', STRONGLY_CONNECTED),
201        ('.', UNBREAKABLE),
202        ('x', DOTTED_NAME),
203        ('==', STRONGLY_CONNECTED),
204        ('37', STRONGLY_CONNECTED),
205        (']', None),
206    ])
207
208  def testFuncCalls(self):
209    code = 'foo(1, 2, 3)\n'
210    tree = self._ParseAndComputePenalties(code)
211    self._CheckPenalties(tree, [
212        ('foo', None),
213        ('(', UNBREAKABLE),
214        ('1', None),
215        (',', UNBREAKABLE),
216        ('2', None),
217        (',', UNBREAKABLE),
218        ('3', None),
219        (')', VERY_STRONGLY_CONNECTED),
220    ])
221
222    # Now a method call, which has more than one trailer
223    code = 'foo.bar.baz(1, 2, 3)\n'
224    tree = self._ParseAndComputePenalties(code)
225    self._CheckPenalties(tree, [
226        ('foo', None),
227        ('.', UNBREAKABLE),
228        ('bar', DOTTED_NAME),
229        ('.', STRONGLY_CONNECTED),
230        ('baz', DOTTED_NAME),
231        ('(', STRONGLY_CONNECTED),
232        ('1', None),
233        (',', UNBREAKABLE),
234        ('2', None),
235        (',', UNBREAKABLE),
236        ('3', None),
237        (')', VERY_STRONGLY_CONNECTED),
238    ])
239
240    # Test single generator argument.
241    code = 'max(i for i in xrange(10))\n'
242    tree = self._ParseAndComputePenalties(code)
243    self._CheckPenalties(tree, [
244        ('max', None),
245        ('(', UNBREAKABLE),
246        ('i', 0),
247        ('for', 0),
248        ('i', STRONGLY_CONNECTED),
249        ('in', STRONGLY_CONNECTED),
250        ('xrange', STRONGLY_CONNECTED),
251        ('(', UNBREAKABLE),
252        ('10', STRONGLY_CONNECTED),
253        (')', VERY_STRONGLY_CONNECTED),
254        (')', VERY_STRONGLY_CONNECTED),
255    ])
256
257
258if __name__ == '__main__':
259  unittest.main()
260