• 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"""Subtype assigner for format tokens.
15
16This module assigns extra type information to format tokens. This information is
17more specific than whether something is an operator or an identifier. For
18instance, it can specify if a node in the tree is part of a subscript.
19
20  AssignSubtypes(): the main function exported by this module.
21
22Annotations:
23  subtype: The subtype of a pytree token. See 'subtypes' module for a list of
24      subtypes.
25"""
26
27from lib2to3 import pytree
28from lib2to3.pgen2 import token as grammar_token
29from lib2to3.pygram import python_symbols as syms
30
31from yapf.yapflib import format_token
32from yapf.yapflib import pytree_utils
33from yapf.yapflib import pytree_visitor
34from yapf.yapflib import style
35from yapf.yapflib import subtypes
36
37
38def AssignSubtypes(tree):
39  """Run the subtype assigner visitor over the tree, modifying it in place.
40
41  Arguments:
42    tree: the top-level pytree node to annotate with subtypes.
43  """
44  subtype_assigner = _SubtypeAssigner()
45  subtype_assigner.Visit(tree)
46
47
48# Map tokens in argument lists to their respective subtype.
49_ARGLIST_TOKEN_TO_SUBTYPE = {
50    '=': subtypes.DEFAULT_OR_NAMED_ASSIGN,
51    ':': subtypes.TYPED_NAME,
52    '*': subtypes.VARARGS_STAR,
53    '**': subtypes.KWARGS_STAR_STAR,
54}
55
56
57class _SubtypeAssigner(pytree_visitor.PyTreeVisitor):
58  """_SubtypeAssigner - see file-level docstring for detailed description.
59
60  The subtype is added as an annotation to the pytree token.
61  """
62
63  def Visit_dictsetmaker(self, node):  # pylint: disable=invalid-name
64    # dictsetmaker ::= (test ':' test (comp_for |
65    #                                   (',' test ':' test)* [','])) |
66    #                  (test (comp_for | (',' test)* [',']))
67    for child in node.children:
68      self.Visit(child)
69
70    comp_for = False
71    dict_maker = False
72
73    for child in node.children:
74      if pytree_utils.NodeName(child) == 'comp_for':
75        comp_for = True
76        _AppendFirstLeafTokenSubtype(child, subtypes.DICT_SET_GENERATOR)
77      elif child.type in (grammar_token.COLON, grammar_token.DOUBLESTAR):
78        dict_maker = True
79
80    if not comp_for and dict_maker:
81      last_was_colon = False
82      unpacking = False
83      for child in node.children:
84        if child.type == grammar_token.DOUBLESTAR:
85          _AppendFirstLeafTokenSubtype(child, subtypes.KWARGS_STAR_STAR)
86        if last_was_colon:
87          if style.Get('INDENT_DICTIONARY_VALUE'):
88            _InsertPseudoParentheses(child)
89          else:
90            _AppendFirstLeafTokenSubtype(child, subtypes.DICTIONARY_VALUE)
91        elif (isinstance(child, pytree.Node) or
92              (not child.value.startswith('#') and child.value not in '{:,')):
93          # Mark the first leaf of a key entry as a DICTIONARY_KEY. We
94          # normally want to split before them if the dictionary cannot exist
95          # on a single line.
96          if not unpacking or pytree_utils.FirstLeafNode(child).value == '**':
97            _AppendFirstLeafTokenSubtype(child, subtypes.DICTIONARY_KEY)
98          _AppendSubtypeRec(child, subtypes.DICTIONARY_KEY_PART)
99        last_was_colon = child.type == grammar_token.COLON
100        if child.type == grammar_token.DOUBLESTAR:
101          unpacking = True
102        elif last_was_colon:
103          unpacking = False
104
105  def Visit_expr_stmt(self, node):  # pylint: disable=invalid-name
106    # expr_stmt ::= testlist_star_expr (augassign (yield_expr|testlist)
107    #               | ('=' (yield_expr|testlist_star_expr))*)
108    for child in node.children:
109      self.Visit(child)
110      if isinstance(child, pytree.Leaf) and child.value == '=':
111        _AppendTokenSubtype(child, subtypes.ASSIGN_OPERATOR)
112
113  def Visit_or_test(self, node):  # pylint: disable=invalid-name
114    # or_test ::= and_test ('or' and_test)*
115    for child in node.children:
116      self.Visit(child)
117      if isinstance(child, pytree.Leaf) and child.value == 'or':
118        _AppendTokenSubtype(child, subtypes.BINARY_OPERATOR)
119
120  def Visit_and_test(self, node):  # pylint: disable=invalid-name
121    # and_test ::= not_test ('and' not_test)*
122    for child in node.children:
123      self.Visit(child)
124      if isinstance(child, pytree.Leaf) and child.value == 'and':
125        _AppendTokenSubtype(child, subtypes.BINARY_OPERATOR)
126
127  def Visit_not_test(self, node):  # pylint: disable=invalid-name
128    # not_test ::= 'not' not_test | comparison
129    for child in node.children:
130      self.Visit(child)
131      if isinstance(child, pytree.Leaf) and child.value == 'not':
132        _AppendTokenSubtype(child, subtypes.UNARY_OPERATOR)
133
134  def Visit_comparison(self, node):  # pylint: disable=invalid-name
135    # comparison ::= expr (comp_op expr)*
136    # comp_op ::= '<'|'>'|'=='|'>='|'<='|'<>'|'!='|'in'|'not in'|'is'|'is not'
137    for child in node.children:
138      self.Visit(child)
139      if (isinstance(child, pytree.Leaf) and
140          child.value in {'<', '>', '==', '>=', '<=', '<>', '!=', 'in', 'is'}):
141        _AppendTokenSubtype(child, subtypes.BINARY_OPERATOR)
142      elif pytree_utils.NodeName(child) == 'comp_op':
143        for grandchild in child.children:
144          _AppendTokenSubtype(grandchild, subtypes.BINARY_OPERATOR)
145
146  def Visit_star_expr(self, node):  # pylint: disable=invalid-name
147    # star_expr ::= '*' expr
148    for child in node.children:
149      self.Visit(child)
150      if isinstance(child, pytree.Leaf) and child.value == '*':
151        _AppendTokenSubtype(child, subtypes.UNARY_OPERATOR)
152        _AppendTokenSubtype(child, subtypes.VARARGS_STAR)
153
154  def Visit_expr(self, node):  # pylint: disable=invalid-name
155    # expr ::= xor_expr ('|' xor_expr)*
156    for child in node.children:
157      self.Visit(child)
158      if isinstance(child, pytree.Leaf) and child.value == '|':
159        _AppendTokenSubtype(child, subtypes.BINARY_OPERATOR)
160
161  def Visit_xor_expr(self, node):  # pylint: disable=invalid-name
162    # xor_expr ::= and_expr ('^' and_expr)*
163    for child in node.children:
164      self.Visit(child)
165      if isinstance(child, pytree.Leaf) and child.value == '^':
166        _AppendTokenSubtype(child, subtypes.BINARY_OPERATOR)
167
168  def Visit_and_expr(self, node):  # pylint: disable=invalid-name
169    # and_expr ::= shift_expr ('&' shift_expr)*
170    for child in node.children:
171      self.Visit(child)
172      if isinstance(child, pytree.Leaf) and child.value == '&':
173        _AppendTokenSubtype(child, subtypes.BINARY_OPERATOR)
174
175  def Visit_shift_expr(self, node):  # pylint: disable=invalid-name
176    # shift_expr ::= arith_expr (('<<'|'>>') arith_expr)*
177    for child in node.children:
178      self.Visit(child)
179      if isinstance(child, pytree.Leaf) and child.value in {'<<', '>>'}:
180        _AppendTokenSubtype(child, subtypes.BINARY_OPERATOR)
181
182  def Visit_arith_expr(self, node):  # pylint: disable=invalid-name
183    # arith_expr ::= term (('+'|'-') term)*
184    for child in node.children:
185      self.Visit(child)
186      if _IsAExprOperator(child):
187        _AppendTokenSubtype(child, subtypes.BINARY_OPERATOR)
188
189    if _IsSimpleExpression(node):
190      for child in node.children:
191        if _IsAExprOperator(child):
192          _AppendTokenSubtype(child, subtypes.SIMPLE_EXPRESSION)
193
194  def Visit_term(self, node):  # pylint: disable=invalid-name
195    # term ::= factor (('*'|'/'|'%'|'//'|'@') factor)*
196    for child in node.children:
197      self.Visit(child)
198      if _IsMExprOperator(child):
199        _AppendTokenSubtype(child, subtypes.BINARY_OPERATOR)
200
201    if _IsSimpleExpression(node):
202      for child in node.children:
203        if _IsMExprOperator(child):
204          _AppendTokenSubtype(child, subtypes.SIMPLE_EXPRESSION)
205
206  def Visit_factor(self, node):  # pylint: disable=invalid-name
207    # factor ::= ('+'|'-'|'~') factor | power
208    for child in node.children:
209      self.Visit(child)
210      if isinstance(child, pytree.Leaf) and child.value in '+-~':
211        _AppendTokenSubtype(child, subtypes.UNARY_OPERATOR)
212
213  def Visit_power(self, node):  # pylint: disable=invalid-name
214    # power ::= atom trailer* ['**' factor]
215    for child in node.children:
216      self.Visit(child)
217      if isinstance(child, pytree.Leaf) and child.value == '**':
218        _AppendTokenSubtype(child, subtypes.BINARY_OPERATOR)
219
220  def Visit_trailer(self, node):  # pylint: disable=invalid-name
221    for child in node.children:
222      self.Visit(child)
223      if isinstance(child, pytree.Leaf) and child.value in '[]':
224        _AppendTokenSubtype(child, subtypes.SUBSCRIPT_BRACKET)
225
226  def Visit_subscript(self, node):  # pylint: disable=invalid-name
227    # subscript ::= test | [test] ':' [test] [sliceop]
228    for child in node.children:
229      self.Visit(child)
230      if isinstance(child, pytree.Leaf) and child.value == ':':
231        _AppendTokenSubtype(child, subtypes.SUBSCRIPT_COLON)
232
233  def Visit_sliceop(self, node):  # pylint: disable=invalid-name
234    # sliceop ::= ':' [test]
235    for child in node.children:
236      self.Visit(child)
237      if isinstance(child, pytree.Leaf) and child.value == ':':
238        _AppendTokenSubtype(child, subtypes.SUBSCRIPT_COLON)
239
240  def Visit_argument(self, node):  # pylint: disable=invalid-name
241    # argument ::=
242    #     test [comp_for] | test '=' test
243    self._ProcessArgLists(node)
244
245  def Visit_arglist(self, node):  # pylint: disable=invalid-name
246    # arglist ::=
247    #     (argument ',')* (argument [',']
248    #                     | '*' test (',' argument)* [',' '**' test]
249    #                     | '**' test)
250    self._ProcessArgLists(node)
251    _SetArgListSubtype(node, subtypes.DEFAULT_OR_NAMED_ASSIGN,
252                       subtypes.DEFAULT_OR_NAMED_ASSIGN_ARG_LIST)
253
254  def Visit_tname(self, node):  # pylint: disable=invalid-name
255    self._ProcessArgLists(node)
256    _SetArgListSubtype(node, subtypes.DEFAULT_OR_NAMED_ASSIGN,
257                       subtypes.DEFAULT_OR_NAMED_ASSIGN_ARG_LIST)
258
259  def Visit_decorator(self, node):  # pylint: disable=invalid-name
260    # decorator ::=
261    #     '@' dotted_name [ '(' [arglist] ')' ] NEWLINE
262    for child in node.children:
263      if isinstance(child, pytree.Leaf) and child.value == '@':
264        _AppendTokenSubtype(child, subtype=subtypes.DECORATOR)
265      self.Visit(child)
266
267  def Visit_funcdef(self, node):  # pylint: disable=invalid-name
268    # funcdef ::=
269    #     'def' NAME parameters ['->' test] ':' suite
270    for child in node.children:
271      if child.type == grammar_token.NAME and child.value != 'def':
272        _AppendTokenSubtype(child, subtypes.FUNC_DEF)
273        break
274    for child in node.children:
275      self.Visit(child)
276
277  def Visit_parameters(self, node):  # pylint: disable=invalid-name
278    # parameters ::= '(' [typedargslist] ')'
279    self._ProcessArgLists(node)
280    if len(node.children) > 2:
281      _AppendFirstLeafTokenSubtype(node.children[1], subtypes.PARAMETER_START)
282      _AppendLastLeafTokenSubtype(node.children[-2], subtypes.PARAMETER_STOP)
283
284  def Visit_typedargslist(self, node):  # pylint: disable=invalid-name
285    # typedargslist ::=
286    #     ((tfpdef ['=' test] ',')*
287    #          ('*' [tname] (',' tname ['=' test])* [',' '**' tname]
288    #           | '**' tname)
289    #     | tfpdef ['=' test] (',' tfpdef ['=' test])* [','])
290    self._ProcessArgLists(node)
291    _SetArgListSubtype(node, subtypes.DEFAULT_OR_NAMED_ASSIGN,
292                       subtypes.DEFAULT_OR_NAMED_ASSIGN_ARG_LIST)
293    tname = False
294    if not node.children:
295      return
296
297    _AppendFirstLeafTokenSubtype(node.children[0], subtypes.PARAMETER_START)
298    _AppendLastLeafTokenSubtype(node.children[-1], subtypes.PARAMETER_STOP)
299
300    tname = pytree_utils.NodeName(node.children[0]) == 'tname'
301    for i in range(1, len(node.children)):
302      prev_child = node.children[i - 1]
303      child = node.children[i]
304      if prev_child.type == grammar_token.COMMA:
305        _AppendFirstLeafTokenSubtype(child, subtypes.PARAMETER_START)
306      elif child.type == grammar_token.COMMA:
307        _AppendLastLeafTokenSubtype(prev_child, subtypes.PARAMETER_STOP)
308
309      if pytree_utils.NodeName(child) == 'tname':
310        tname = True
311        _SetArgListSubtype(child, subtypes.TYPED_NAME,
312                           subtypes.TYPED_NAME_ARG_LIST)
313      elif child.type == grammar_token.COMMA:
314        tname = False
315      elif child.type == grammar_token.EQUAL and tname:
316        _AppendTokenSubtype(child, subtype=subtypes.TYPED_NAME)
317        tname = False
318
319  def Visit_varargslist(self, node):  # pylint: disable=invalid-name
320    # varargslist ::=
321    #     ((vfpdef ['=' test] ',')*
322    #          ('*' [vname] (',' vname ['=' test])*  [',' '**' vname]
323    #           | '**' vname)
324    #      | vfpdef ['=' test] (',' vfpdef ['=' test])* [','])
325    self._ProcessArgLists(node)
326    for child in node.children:
327      self.Visit(child)
328      if isinstance(child, pytree.Leaf) and child.value == '=':
329        _AppendTokenSubtype(child, subtypes.VARARGS_LIST)
330
331  def Visit_comp_for(self, node):  # pylint: disable=invalid-name
332    # comp_for ::= 'for' exprlist 'in' testlist_safe [comp_iter]
333    _AppendSubtypeRec(node, subtypes.COMP_FOR)
334    # Mark the previous node as COMP_EXPR unless this is a nested comprehension
335    # as these will have the outer comprehension as their previous node.
336    attr = pytree_utils.GetNodeAnnotation(node.parent,
337                                          pytree_utils.Annotation.SUBTYPE)
338    if not attr or subtypes.COMP_FOR not in attr:
339      _AppendSubtypeRec(node.parent.children[0], subtypes.COMP_EXPR)
340    self.DefaultNodeVisit(node)
341
342  def Visit_old_comp_for(self, node):  # pylint: disable=invalid-name
343    # Python 3.7
344    self.Visit_comp_for(node)
345
346  def Visit_comp_if(self, node):  # pylint: disable=invalid-name
347    # comp_if ::= 'if' old_test [comp_iter]
348    _AppendSubtypeRec(node, subtypes.COMP_IF)
349    self.DefaultNodeVisit(node)
350
351  def Visit_old_comp_if(self, node):  # pylint: disable=invalid-name
352    # Python 3.7
353    self.Visit_comp_if(node)
354
355  def _ProcessArgLists(self, node):
356    """Common method for processing argument lists."""
357    for child in node.children:
358      self.Visit(child)
359      if isinstance(child, pytree.Leaf):
360        _AppendTokenSubtype(
361            child,
362            subtype=_ARGLIST_TOKEN_TO_SUBTYPE.get(child.value, subtypes.NONE))
363
364
365def _SetArgListSubtype(node, node_subtype, list_subtype):
366  """Set named assign subtype on elements in a arg list."""
367
368  def HasSubtype(node):
369    """Return True if the arg list has a named assign subtype."""
370    if isinstance(node, pytree.Leaf):
371      return node_subtype in pytree_utils.GetNodeAnnotation(
372          node, pytree_utils.Annotation.SUBTYPE, set())
373
374    for child in node.children:
375      node_name = pytree_utils.NodeName(child)
376      if node_name not in {'atom', 'arglist', 'power'}:
377        if HasSubtype(child):
378          return True
379
380    return False
381
382  if not HasSubtype(node):
383    return
384
385  for child in node.children:
386    node_name = pytree_utils.NodeName(child)
387    if node_name not in {'atom', 'COMMA'}:
388      _AppendFirstLeafTokenSubtype(child, list_subtype)
389
390
391def _AppendTokenSubtype(node, subtype):
392  """Append the token's subtype only if it's not already set."""
393  pytree_utils.AppendNodeAnnotation(node, pytree_utils.Annotation.SUBTYPE,
394                                    subtype)
395
396
397def _AppendFirstLeafTokenSubtype(node, subtype):
398  """Append the first leaf token's subtypes."""
399  if isinstance(node, pytree.Leaf):
400    _AppendTokenSubtype(node, subtype)
401    return
402  _AppendFirstLeafTokenSubtype(node.children[0], subtype)
403
404
405def _AppendLastLeafTokenSubtype(node, subtype):
406  """Append the last leaf token's subtypes."""
407  if isinstance(node, pytree.Leaf):
408    _AppendTokenSubtype(node, subtype)
409    return
410  _AppendLastLeafTokenSubtype(node.children[-1], subtype)
411
412
413def _AppendSubtypeRec(node, subtype, force=True):
414  """Append the leafs in the node to the given subtype."""
415  if isinstance(node, pytree.Leaf):
416    _AppendTokenSubtype(node, subtype)
417    return
418  for child in node.children:
419    _AppendSubtypeRec(child, subtype, force=force)
420
421
422def _InsertPseudoParentheses(node):
423  """Insert pseudo parentheses so that dicts can be formatted correctly."""
424  comment_node = None
425  if isinstance(node, pytree.Node):
426    if node.children[-1].type == grammar_token.COMMENT:
427      comment_node = node.children[-1].clone()
428      node.children[-1].remove()
429
430  first = pytree_utils.FirstLeafNode(node)
431  last = pytree_utils.LastLeafNode(node)
432
433  if first == last and first.type == grammar_token.COMMENT:
434    # A comment was inserted before the value, which is a pytree.Leaf.
435    # Encompass the dictionary's value into an ATOM node.
436    last = first.next_sibling
437    last_clone = last.clone()
438    new_node = pytree.Node(syms.atom, [first.clone(), last_clone])
439    for orig_leaf, clone_leaf in zip(last.leaves(), last_clone.leaves()):
440      pytree_utils.CopyYapfAnnotations(orig_leaf, clone_leaf)
441      if hasattr(orig_leaf, 'is_pseudo'):
442        clone_leaf.is_pseudo = orig_leaf.is_pseudo
443
444    node.replace(new_node)
445    node = new_node
446    last.remove()
447
448    first = pytree_utils.FirstLeafNode(node)
449    last = pytree_utils.LastLeafNode(node)
450
451  lparen = pytree.Leaf(
452      grammar_token.LPAR,
453      u'(',
454      context=('', (first.get_lineno(), first.column - 1)))
455  last_lineno = last.get_lineno()
456  if last.type == grammar_token.STRING and '\n' in last.value:
457    last_lineno += last.value.count('\n')
458
459  if last.type == grammar_token.STRING and '\n' in last.value:
460    last_column = len(last.value.split('\n')[-1]) + 1
461  else:
462    last_column = last.column + len(last.value) + 1
463  rparen = pytree.Leaf(
464      grammar_token.RPAR, u')', context=('', (last_lineno, last_column)))
465
466  lparen.is_pseudo = True
467  rparen.is_pseudo = True
468
469  if isinstance(node, pytree.Node):
470    node.insert_child(0, lparen)
471    node.append_child(rparen)
472    if comment_node:
473      node.append_child(comment_node)
474    _AppendFirstLeafTokenSubtype(node, subtypes.DICTIONARY_VALUE)
475  else:
476    clone = node.clone()
477    for orig_leaf, clone_leaf in zip(node.leaves(), clone.leaves()):
478      pytree_utils.CopyYapfAnnotations(orig_leaf, clone_leaf)
479    new_node = pytree.Node(syms.atom, [lparen, clone, rparen])
480    node.replace(new_node)
481    _AppendFirstLeafTokenSubtype(clone, subtypes.DICTIONARY_VALUE)
482
483
484def _IsAExprOperator(node):
485  return isinstance(node, pytree.Leaf) and node.value in {'+', '-'}
486
487
488def _IsMExprOperator(node):
489  return isinstance(node,
490                    pytree.Leaf) and node.value in {'*', '/', '%', '//', '@'}
491
492
493def _IsSimpleExpression(node):
494  """A node with only leafs as children."""
495  return all(isinstance(child, pytree.Leaf) for child in node.children)
496