djavu-asus/elpy/rpc-venv/lib/python3.11/site-packages/yapf/pytree/split_penalty.py

   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 """Computation of split penalties before/between tokens."""
  15
  16 import re
  17
  18 from yapf_third_party._ylib2to3 import pytree
  19 from yapf_third_party._ylib2to3.pgen2 import token as grammar_token
  20
  21 from yapf.pytree import pytree_utils
  22 from yapf.pytree import pytree_visitor
  23 from yapf.yapflib import style
  24 from yapf.yapflib import subtypes
  25
  26 # TODO(morbo): Document the annotations in a centralized place. E.g., the
  27 # README file.
  28 UNBREAKABLE = 1000 * 1000
  29 NAMED_ASSIGN = 15000
  30 DOTTED_NAME = 4000
  31 VERY_STRONGLY_CONNECTED = 3500
  32 STRONGLY_CONNECTED = 3000
  33 CONNECTED = 500
  34 TOGETHER = 100
  35
  36 OR_TEST = 1000
  37 AND_TEST = 1100
  38 NOT_TEST = 1200
  39 COMPARISON = 1300
  40 STAR_EXPR = 1300
  41 EXPR = 1400
  42 XOR_EXPR = 1500
  43 AND_EXPR = 1700
  44 SHIFT_EXPR = 1800
  45 ARITH_EXPR = 1900
  46 TERM = 2000
  47 FACTOR = 2100
  48 POWER = 2200
  49 ATOM = 2300
  50 ONE_ELEMENT_ARGUMENT = 500
  51 SUBSCRIPT = 6000
  52
  53
  54 def ComputeSplitPenalties(tree):
  55   """Compute split penalties on tokens in the given parse tree.
  56
  57   Arguments:
  58     tree: the top-level pytree node to annotate with penalties.
  59   """
  60   _SplitPenaltyAssigner().Visit(tree)
  61
  62
  63 class _SplitPenaltyAssigner(pytree_visitor.PyTreeVisitor):
  64   """Assigns split penalties to tokens, based on parse tree structure.
  65
  66   Split penalties are attached as annotations to tokens.
  67   """
  68
  69   def Visit(self, node):
  70     if not hasattr(node, 'is_pseudo'):  # Ignore pseudo tokens.
  71       super(_SplitPenaltyAssigner, self).Visit(node)
  72
  73   def Visit_import_as_names(self, node):  # pyline: disable=invalid-name
  74     # import_as_names ::= import_as_name (',' import_as_name)* [',']
  75     self.DefaultNodeVisit(node)
  76     prev_child = None
  77     for child in node.children:
  78       if (prev_child and isinstance(prev_child, pytree.Leaf) and
  79           prev_child.value == ','):
  80         _SetSplitPenalty(child, style.Get('SPLIT_PENALTY_IMPORT_NAMES'))
  81       prev_child = child
  82
  83   def Visit_classdef(self, node):  # pylint: disable=invalid-name
  84     # classdef ::= 'class' NAME ['(' [arglist] ')'] ':' suite
  85     #
  86     # NAME
  87     _SetUnbreakable(node.children[1])
  88     if len(node.children) > 4:
  89       # opening '('
  90       _SetUnbreakable(node.children[2])
  91     # ':'
  92     _SetUnbreakable(node.children[-2])
  93     self.DefaultNodeVisit(node)
  94
  95   def Visit_funcdef(self, node):  # pylint: disable=invalid-name
  96     # funcdef ::= 'def' NAME parameters ['->' test] ':' suite
  97     #
  98     # Can't break before the function name and before the colon. The parameters
  99     # are handled by child iteration.
 100     colon_idx = 1
 101     while pytree_utils.NodeName(node.children[colon_idx]) == 'simple_stmt':
 102       colon_idx += 1
 103     _SetUnbreakable(node.children[colon_idx])
 104     arrow_idx = -1
 105     while colon_idx < len(node.children):
 106       if isinstance(node.children[colon_idx], pytree.Leaf):
 107         if node.children[colon_idx].value == ':':
 108           break
 109         if node.children[colon_idx].value == '->':
 110           arrow_idx = colon_idx
 111       colon_idx += 1
 112     _SetUnbreakable(node.children[colon_idx])
 113     self.DefaultNodeVisit(node)
 114     if arrow_idx > 0:
 115       _SetSplitPenalty(
 116           pytree_utils.LastLeafNode(node.children[arrow_idx - 1]), 0)
 117       _SetUnbreakable(node.children[arrow_idx])
 118       _SetStronglyConnected(node.children[arrow_idx + 1])
 119
 120   def Visit_lambdef(self, node):  # pylint: disable=invalid-name
 121     # lambdef ::= 'lambda' [varargslist] ':' test
 122     # Loop over the lambda up to and including the colon.
 123     allow_multiline_lambdas = style.Get('ALLOW_MULTILINE_LAMBDAS')
 124     if not allow_multiline_lambdas:
 125       for child in node.children:
 126         if child.type == grammar_token.COMMENT:
 127           if re.search(r'pylint:.*disable=.*\bg-long-lambda', child.value):
 128             allow_multiline_lambdas = True
 129             break
 130
 131     if allow_multiline_lambdas:
 132       _SetExpressionPenalty(node, STRONGLY_CONNECTED)
 133     else:
 134       _SetExpressionPenalty(node, VERY_STRONGLY_CONNECTED)
 135
 136   def Visit_parameters(self, node):  # pylint: disable=invalid-name
 137     # parameters ::= '(' [typedargslist] ')'
 138     self.DefaultNodeVisit(node)
 139
 140     # Can't break before the opening paren of a parameter list.
 141     _SetUnbreakable(node.children[0])
 142     if not (style.Get('INDENT_CLOSING_BRACKETS') or
 143             style.Get('DEDENT_CLOSING_BRACKETS')):
 144       _SetStronglyConnected(node.children[-1])
 145
 146   def Visit_arglist(self, node):  # pylint: disable=invalid-name
 147     # arglist ::= argument (',' argument)* [',']
 148     if node.children[0].type == grammar_token.STAR:
 149       # Python 3 treats a star expression as a specific expression type.
 150       # Process it in that method.
 151       self.Visit_star_expr(node)
 152       return
 153
 154     self.DefaultNodeVisit(node)
 155
 156     for index in range(1, len(node.children)):
 157       child = node.children[index]
 158       if isinstance(child, pytree.Leaf) and child.value == ',':
 159         _SetUnbreakable(child)
 160
 161     for child in node.children:
 162       if pytree_utils.NodeName(child) == 'atom':
 163         _IncreasePenalty(child, CONNECTED)
 164
 165   def Visit_argument(self, node):  # pylint: disable=invalid-name
 166     # argument ::= test [comp_for] | test '=' test  # Really [keyword '='] test
 167     self.DefaultNodeVisit(node)
 168
 169     for index in range(1, len(node.children) - 1):
 170       child = node.children[index]
 171       if isinstance(child, pytree.Leaf) and child.value == '=':
 172         _SetSplitPenalty(
 173             pytree_utils.FirstLeafNode(node.children[index]), NAMED_ASSIGN)
 174         _SetSplitPenalty(
 175             pytree_utils.FirstLeafNode(node.children[index + 1]), NAMED_ASSIGN)
 176
 177   def Visit_tname(self, node):  # pylint: disable=invalid-name
 178     # tname ::= NAME [':' test]
 179     self.DefaultNodeVisit(node)
 180
 181     for index in range(1, len(node.children) - 1):
 182       child = node.children[index]
 183       if isinstance(child, pytree.Leaf) and child.value == ':':
 184         _SetSplitPenalty(
 185             pytree_utils.FirstLeafNode(node.children[index]), NAMED_ASSIGN)
 186         _SetSplitPenalty(
 187             pytree_utils.FirstLeafNode(node.children[index + 1]), NAMED_ASSIGN)
 188
 189   def Visit_dotted_name(self, node):  # pylint: disable=invalid-name
 190     # dotted_name ::= NAME ('.' NAME)*
 191     for child in node.children:
 192       self.Visit(child)
 193     start = 2 if hasattr(node.children[0], 'is_pseudo') else 1
 194     for i in range(start, len(node.children)):
 195       _SetUnbreakable(node.children[i])
 196
 197   def Visit_dictsetmaker(self, node):  # pylint: disable=invalid-name
 198     # dictsetmaker ::= ( (test ':' test
 199     #                      (comp_for | (',' test ':' test)* [','])) |
 200     #                    (test (comp_for | (',' test)* [','])) )
 201     for child in node.children:
 202       self.Visit(child)
 203       if child.type == grammar_token.COLON:
 204         # This is a key to a dictionary. We don't want to split the key if at
 205         # all possible.
 206         _SetStronglyConnected(child)
 207
 208   def Visit_trailer(self, node):  # pylint: disable=invalid-name
 209     # trailer ::= '(' [arglist] ')' | '[' subscriptlist ']' | '.' NAME
 210     if node.children[0].value == '.':
 211       before = style.Get('SPLIT_BEFORE_DOT')
 212       _SetSplitPenalty(node.children[0],
 213                        VERY_STRONGLY_CONNECTED if before else DOTTED_NAME)
 214       _SetSplitPenalty(node.children[1],
 215                        DOTTED_NAME if before else VERY_STRONGLY_CONNECTED)
 216     elif len(node.children) == 2:
 217       # Don't split an empty argument list if at all possible.
 218       _SetSplitPenalty(node.children[1], VERY_STRONGLY_CONNECTED)
 219     elif len(node.children) == 3:
 220       name = pytree_utils.NodeName(node.children[1])
 221       if name in {'argument', 'comparison'}:
 222         # Don't split an argument list with one element if at all possible.
 223         _SetStronglyConnected(node.children[1])
 224         if (len(node.children[1].children) > 1 and
 225             pytree_utils.NodeName(node.children[1].children[1]) == 'comp_for'):
 226           # Don't penalize splitting before a comp_for expression.
 227           _SetSplitPenalty(pytree_utils.FirstLeafNode(node.children[1]), 0)
 228         else:
 229           _SetSplitPenalty(
 230               pytree_utils.FirstLeafNode(node.children[1]),
 231               ONE_ELEMENT_ARGUMENT)
 232       elif (node.children[0].type == grammar_token.LSQB and
 233             len(node.children[1].children) > 2 and
 234             (name.endswith('_test') or name.endswith('_expr'))):
 235         _SetStronglyConnected(node.children[1].children[0])
 236         _SetStronglyConnected(node.children[1].children[2])
 237
 238         # Still allow splitting around the operator.
 239         split_before = ((name.endswith('_test') and
 240                          style.Get('SPLIT_BEFORE_LOGICAL_OPERATOR')) or
 241                         (name.endswith('_expr') and
 242                          style.Get('SPLIT_BEFORE_BITWISE_OPERATOR')))
 243         if split_before:
 244           _SetSplitPenalty(
 245               pytree_utils.LastLeafNode(node.children[1].children[1]), 0)
 246         else:
 247           _SetSplitPenalty(
 248               pytree_utils.FirstLeafNode(node.children[1].children[2]), 0)
 249
 250         # Don't split the ending bracket of a subscript list.
 251         _RecAnnotate(node.children[-1], pytree_utils.Annotation.SPLIT_PENALTY,
 252                      VERY_STRONGLY_CONNECTED)
 253       elif name not in {
 254           'arglist', 'argument', 'term', 'or_test', 'and_test', 'comparison',
 255           'atom', 'power'
 256       }:
 257         # Don't split an argument list with one element if at all possible.
 258         stypes = pytree_utils.GetNodeAnnotation(
 259             pytree_utils.FirstLeafNode(node), pytree_utils.Annotation.SUBTYPE)
 260         if stypes and subtypes.SUBSCRIPT_BRACKET in stypes:
 261           _IncreasePenalty(node, SUBSCRIPT)
 262
 263           # Bump up the split penalty for the first part of a subscript. We
 264           # would rather not split there.
 265           _IncreasePenalty(node.children[1], CONNECTED)
 266         else:
 267           _SetStronglyConnected(node.children[1], node.children[2])
 268
 269       if name == 'arglist':
 270         _SetStronglyConnected(node.children[-1])
 271
 272     self.DefaultNodeVisit(node)
 273
 274   def Visit_power(self, node):  # pylint: disable=invalid-name,missing-docstring
 275     # power ::= atom trailer* ['**' factor]
 276     self.DefaultNodeVisit(node)
 277
 278     # When atom is followed by a trailer, we can not break between them.
 279     # E.g. arr[idx] - no break allowed between 'arr' and '['.
 280     if (len(node.children) > 1 and
 281         pytree_utils.NodeName(node.children[1]) == 'trailer'):
 282       # children[1] itself is a whole trailer: we don't want to
 283       # mark all of it as unbreakable, only its first token: (, [ or .
 284       first = pytree_utils.FirstLeafNode(node.children[1])
 285       if first.value != '.':
 286         _SetUnbreakable(node.children[1].children[0])
 287
 288       # A special case when there are more trailers in the sequence. Given:
 289       #   atom tr1 tr2
 290       # The last token of tr1 and the first token of tr2 comprise an unbreakable
 291       # region. For example: foo.bar.baz(1)
 292       # We can't put breaks between either of the '.', '(', or '[' and the names
 293       # *preceding* them.
 294       prev_trailer_idx = 1
 295       while prev_trailer_idx < len(node.children) - 1:
 296         cur_trailer_idx = prev_trailer_idx + 1
 297         cur_trailer = node.children[cur_trailer_idx]
 298         if pytree_utils.NodeName(cur_trailer) != 'trailer':
 299           break
 300
 301         # Now we know we have two trailers one after the other
 302         prev_trailer = node.children[prev_trailer_idx]
 303         if prev_trailer.children[-1].value != ')':
 304           # Set the previous node unbreakable if it's not a function call:
 305           #   atom tr1() tr2
 306           # It may be necessary (though undesirable) to split up a previous
 307           # function call's parentheses to the next line.
 308           _SetStronglyConnected(prev_trailer.children[-1])
 309         _SetStronglyConnected(cur_trailer.children[0])
 310         prev_trailer_idx = cur_trailer_idx
 311
 312     # We don't want to split before the last ')' of a function call. This also
 313     # takes care of the special case of:
 314     #   atom tr1 tr2 ... trn
 315     # where the 'tr#' are trailers that may end in a ')'.
 316     for trailer in node.children[1:]:
 317       if pytree_utils.NodeName(trailer) != 'trailer':
 318         break
 319       if trailer.children[0].value in '([':
 320         if len(trailer.children) > 2:
 321           stypes = pytree_utils.GetNodeAnnotation(
 322               trailer.children[0], pytree_utils.Annotation.SUBTYPE)
 323           if stypes and subtypes.SUBSCRIPT_BRACKET in stypes:
 324             _SetStronglyConnected(
 325                 pytree_utils.FirstLeafNode(trailer.children[1]))
 326
 327           last_child_node = pytree_utils.LastLeafNode(trailer)
 328           if last_child_node.value.strip().startswith('#'):
 329             last_child_node = last_child_node.prev_sibling
 330           if not (style.Get('INDENT_CLOSING_BRACKETS') or
 331                   style.Get('DEDENT_CLOSING_BRACKETS')):
 332             last = pytree_utils.LastLeafNode(last_child_node.prev_sibling)
 333             if last.value != ',':
 334               if last_child_node.value == ']':
 335                 _SetUnbreakable(last_child_node)
 336               else:
 337                 _SetSplitPenalty(last_child_node, VERY_STRONGLY_CONNECTED)
 338         else:
 339           # If the trailer's children are '()', then make it a strongly
 340           # connected region.  It's sometimes necessary, though undesirable, to
 341           # split the two.
 342           _SetStronglyConnected(trailer.children[-1])
 343
 344   def Visit_subscriptlist(self, node):  # pylint: disable=invalid-name
 345     # subscriptlist ::= subscript (',' subscript)* [',']
 346     self.DefaultNodeVisit(node)
 347     _SetSplitPenalty(pytree_utils.FirstLeafNode(node), 0)
 348     prev_child = None
 349     for child in node.children:
 350       if prev_child and prev_child.type == grammar_token.COMMA:
 351         _SetSplitPenalty(pytree_utils.FirstLeafNode(child), 0)
 352       prev_child = child
 353
 354   def Visit_subscript(self, node):  # pylint: disable=invalid-name
 355     # subscript ::= test | [test] ':' [test] [sliceop]
 356     _SetStronglyConnected(*node.children)
 357     self.DefaultNodeVisit(node)
 358
 359   def Visit_comp_for(self, node):  # pylint: disable=invalid-name
 360     # comp_for ::= 'for' exprlist 'in' testlist_safe [comp_iter]
 361     _SetSplitPenalty(pytree_utils.FirstLeafNode(node), 0)
 362     _SetStronglyConnected(*node.children[1:])
 363     self.DefaultNodeVisit(node)
 364
 365   def Visit_old_comp_for(self, node):  # pylint: disable=invalid-name
 366     # Python 3.7
 367     self.Visit_comp_for(node)
 368
 369   def Visit_comp_if(self, node):  # pylint: disable=invalid-name
 370     # comp_if ::= 'if' old_test [comp_iter]
 371     _SetSplitPenalty(node.children[0],
 372                      style.Get('SPLIT_PENALTY_BEFORE_IF_EXPR'))
 373     _SetStronglyConnected(*node.children[1:])
 374     self.DefaultNodeVisit(node)
 375
 376   def Visit_old_comp_if(self, node):  # pylint: disable=invalid-name
 377     # Python 3.7
 378     self.Visit_comp_if(node)
 379
 380   def Visit_test(self, node):  # pylint: disable=invalid-name
 381     # test ::= or_test ['if' or_test 'else' test] | lambdef
 382     _IncreasePenalty(node, OR_TEST)
 383     self.DefaultNodeVisit(node)
 384
 385   def Visit_or_test(self, node):  # pylint: disable=invalid-name
 386     # or_test ::= and_test ('or' and_test)*
 387     self.DefaultNodeVisit(node)
 388     _IncreasePenalty(node, OR_TEST)
 389     index = 1
 390     while index + 1 < len(node.children):
 391       if style.Get('SPLIT_BEFORE_LOGICAL_OPERATOR'):
 392         _DecrementSplitPenalty(
 393             pytree_utils.FirstLeafNode(node.children[index]), OR_TEST)
 394       else:
 395         _DecrementSplitPenalty(
 396             pytree_utils.FirstLeafNode(node.children[index + 1]), OR_TEST)
 397       index += 2
 398
 399   def Visit_and_test(self, node):  # pylint: disable=invalid-name
 400     # and_test ::= not_test ('and' not_test)*
 401     self.DefaultNodeVisit(node)
 402     _IncreasePenalty(node, AND_TEST)
 403     index = 1
 404     while index + 1 < len(node.children):
 405       if style.Get('SPLIT_BEFORE_LOGICAL_OPERATOR'):
 406         _DecrementSplitPenalty(
 407             pytree_utils.FirstLeafNode(node.children[index]), AND_TEST)
 408       else:
 409         _DecrementSplitPenalty(
 410             pytree_utils.FirstLeafNode(node.children[index + 1]), AND_TEST)
 411       index += 2
 412
 413   def Visit_not_test(self, node):  # pylint: disable=invalid-name
 414     # not_test ::= 'not' not_test | comparison
 415     self.DefaultNodeVisit(node)
 416     _IncreasePenalty(node, NOT_TEST)
 417
 418   def Visit_comparison(self, node):  # pylint: disable=invalid-name
 419     # comparison ::= expr (comp_op expr)*
 420     self.DefaultNodeVisit(node)
 421     if len(node.children) == 3 and _StronglyConnectedCompOp(node):
 422       _IncreasePenalty(node.children[1], VERY_STRONGLY_CONNECTED)
 423       _SetSplitPenalty(
 424           pytree_utils.FirstLeafNode(node.children[2]), STRONGLY_CONNECTED)
 425     else:
 426       _IncreasePenalty(node, COMPARISON)
 427
 428   def Visit_star_expr(self, node):  # pylint: disable=invalid-name
 429     # star_expr ::= '*' expr
 430     self.DefaultNodeVisit(node)
 431     _IncreasePenalty(node, STAR_EXPR)
 432
 433   def Visit_expr(self, node):  # pylint: disable=invalid-name
 434     # expr ::= xor_expr ('|' xor_expr)*
 435     self.DefaultNodeVisit(node)
 436     _IncreasePenalty(node, EXPR)
 437     _SetBitwiseOperandPenalty(node, '|')
 438
 439   def Visit_xor_expr(self, node):  # pylint: disable=invalid-name
 440     # xor_expr ::= and_expr ('^' and_expr)*
 441     self.DefaultNodeVisit(node)
 442     _IncreasePenalty(node, XOR_EXPR)
 443     _SetBitwiseOperandPenalty(node, '^')
 444
 445   def Visit_and_expr(self, node):  # pylint: disable=invalid-name
 446     # and_expr ::= shift_expr ('&' shift_expr)*
 447     self.DefaultNodeVisit(node)
 448     _IncreasePenalty(node, AND_EXPR)
 449     _SetBitwiseOperandPenalty(node, '&')
 450
 451   def Visit_shift_expr(self, node):  # pylint: disable=invalid-name
 452     # shift_expr ::= arith_expr (('<<'|'>>') arith_expr)*
 453     self.DefaultNodeVisit(node)
 454     _IncreasePenalty(node, SHIFT_EXPR)
 455
 456   _ARITH_OPS = frozenset({'PLUS', 'MINUS'})
 457
 458   def Visit_arith_expr(self, node):  # pylint: disable=invalid-name
 459     # arith_expr ::= term (('+'|'-') term)*
 460     self.DefaultNodeVisit(node)
 461     _IncreasePenalty(node, ARITH_EXPR)
 462     _SetExpressionOperandPenalty(node, self._ARITH_OPS)
 463
 464   _TERM_OPS = frozenset({'STAR', 'AT', 'SLASH', 'PERCENT', 'DOUBLESLASH'})
 465
 466   def Visit_term(self, node):  # pylint: disable=invalid-name
 467     # term ::= factor (('*'|'@'|'/'|'%'|'//') factor)*
 468     self.DefaultNodeVisit(node)
 469     _IncreasePenalty(node, TERM)
 470     _SetExpressionOperandPenalty(node, self._TERM_OPS)
 471
 472   def Visit_factor(self, node):  # pyline: disable=invalid-name
 473     # factor ::= ('+'|'-'|'~') factor | power
 474     self.DefaultNodeVisit(node)
 475     _IncreasePenalty(node, FACTOR)
 476
 477   def Visit_atom(self, node):  # pylint: disable=invalid-name
 478     # atom ::= ('(' [yield_expr|testlist_gexp] ')'
 479     #           '[' [listmaker] ']' |
 480     #           '{' [dictsetmaker] '}')
 481     self.DefaultNodeVisit(node)
 482     if (node.children[0].value == '(' and
 483         not hasattr(node.children[0], 'is_pseudo')):
 484       if node.children[-1].value == ')':
 485         if pytree_utils.NodeName(node.parent) == 'if_stmt':
 486           _SetSplitPenalty(node.children[-1], STRONGLY_CONNECTED)
 487         else:
 488           if len(node.children) > 2:
 489             _SetSplitPenalty(pytree_utils.FirstLeafNode(node.children[1]), EXPR)
 490           _SetSplitPenalty(node.children[-1], ATOM)
 491     elif node.children[0].value in '[{' and len(node.children) == 2:
 492       # Keep empty containers together if we can.
 493       _SetUnbreakable(node.children[-1])
 494
 495   def Visit_testlist_gexp(self, node):  # pylint: disable=invalid-name
 496     self.DefaultNodeVisit(node)
 497     prev_was_comma = False
 498     for child in node.children:
 499       if isinstance(child, pytree.Leaf) and child.value == ',':
 500         _SetUnbreakable(child)
 501         prev_was_comma = True
 502       else:
 503         if prev_was_comma:
 504           _SetSplitPenalty(pytree_utils.FirstLeafNode(child), TOGETHER)
 505         prev_was_comma = False
 506
 507
 508 def _SetUnbreakable(node):
 509   """Set an UNBREAKABLE penalty annotation for the given node."""
 510   _RecAnnotate(node, pytree_utils.Annotation.SPLIT_PENALTY, UNBREAKABLE)
 511
 512
 513 def _SetStronglyConnected(*nodes):
 514   """Set a STRONGLY_CONNECTED penalty annotation for the given nodes."""
 515   for node in nodes:
 516     _RecAnnotate(node, pytree_utils.Annotation.SPLIT_PENALTY,
 517                  STRONGLY_CONNECTED)
 518
 519
 520 def _SetExpressionPenalty(node, penalty):
 521   """Set a penalty annotation on children nodes."""
 522
 523   def RecExpression(node, first_child_leaf):
 524     if node is first_child_leaf:
 525       return
 526
 527     if isinstance(node, pytree.Leaf):
 528       if node.value in {'(', 'for', 'if'}:
 529         return
 530       penalty_annotation = pytree_utils.GetNodeAnnotation(
 531           node, pytree_utils.Annotation.SPLIT_PENALTY, default=0)
 532       if penalty_annotation < penalty:
 533         _SetSplitPenalty(node, penalty)
 534     else:
 535       for child in node.children:
 536         RecExpression(child, first_child_leaf)
 537
 538   RecExpression(node, pytree_utils.FirstLeafNode(node))
 539
 540
 541 def _SetBitwiseOperandPenalty(node, op):
 542   for index in range(1, len(node.children) - 1):
 543     child = node.children[index]
 544     if isinstance(child, pytree.Leaf) and child.value == op:
 545       if style.Get('SPLIT_BEFORE_BITWISE_OPERATOR'):
 546         _SetSplitPenalty(child, style.Get('SPLIT_PENALTY_BITWISE_OPERATOR'))
 547       else:
 548         _SetSplitPenalty(
 549             pytree_utils.FirstLeafNode(node.children[index + 1]),
 550             style.Get('SPLIT_PENALTY_BITWISE_OPERATOR'))
 551
 552
 553 def _SetExpressionOperandPenalty(node, ops):
 554   for index in range(1, len(node.children) - 1):
 555     child = node.children[index]
 556     if pytree_utils.NodeName(child) in ops:
 557       if style.Get('SPLIT_BEFORE_ARITHMETIC_OPERATOR'):
 558         _SetSplitPenalty(child, style.Get('SPLIT_PENALTY_ARITHMETIC_OPERATOR'))
 559       else:
 560         _SetSplitPenalty(
 561             pytree_utils.FirstLeafNode(node.children[index + 1]),
 562             style.Get('SPLIT_PENALTY_ARITHMETIC_OPERATOR'))
 563
 564
 565 def _IncreasePenalty(node, amt):
 566   """Increase a penalty annotation on children nodes."""
 567
 568   def RecExpression(node, first_child_leaf):
 569     if node is first_child_leaf:
 570       return
 571
 572     if isinstance(node, pytree.Leaf):
 573       if node.value in {'(', 'for'}:
 574         return
 575       penalty = pytree_utils.GetNodeAnnotation(
 576           node, pytree_utils.Annotation.SPLIT_PENALTY, default=0)
 577       _SetSplitPenalty(node, penalty + amt)
 578     else:
 579       for child in node.children:
 580         RecExpression(child, first_child_leaf)
 581
 582   RecExpression(node, pytree_utils.FirstLeafNode(node))
 583
 584
 585 def _RecAnnotate(tree, annotate_name, annotate_value):
 586   """Recursively set the given annotation on all leafs of the subtree.
 587
 588   Takes care to only increase the penalty. If the node already has a higher
 589   or equal penalty associated with it, this is a no-op.
 590
 591   Args:
 592     tree: subtree to annotate
 593     annotate_name: name of the annotation to set
 594     annotate_value: value of the annotation to set
 595   """
 596   for child in tree.children:
 597     _RecAnnotate(child, annotate_name, annotate_value)
 598   if isinstance(tree, pytree.Leaf):
 599     cur_annotate = pytree_utils.GetNodeAnnotation(
 600         tree, annotate_name, default=0)
 601     if cur_annotate < annotate_value:
 602       pytree_utils.SetNodeAnnotation(tree, annotate_name, annotate_value)
 603
 604
 605 _COMP_OPS = frozenset({'==', '!=', '<=', '<', '>', '>=', '<>', 'in', 'is'})
 606
 607
 608 def _StronglyConnectedCompOp(op):
 609   if (len(op.children[1].children) == 2 and
 610       pytree_utils.NodeName(op.children[1]) == 'comp_op'):
 611     if (pytree_utils.FirstLeafNode(op.children[1]).value == 'not' and
 612         pytree_utils.LastLeafNode(op.children[1]).value == 'in'):
 613       return True
 614     if (pytree_utils.FirstLeafNode(op.children[1]).value == 'is' and
 615         pytree_utils.LastLeafNode(op.children[1]).value == 'not'):
 616       return True
 617   if (isinstance(op.children[1], pytree.Leaf) and
 618       op.children[1].value in _COMP_OPS):
 619     return True
 620   return False
 621
 622
 623 def _DecrementSplitPenalty(node, amt):
 624   penalty = pytree_utils.GetNodeAnnotation(
 625       node, pytree_utils.Annotation.SPLIT_PENALTY, default=amt)
 626   penalty = penalty - amt if amt < penalty else 0
 627   _SetSplitPenalty(node, penalty)
 628
 629
 630 def _SetSplitPenalty(node, penalty):
 631   pytree_utils.SetNodeAnnotation(node, pytree_utils.Annotation.SPLIT_PENALTY,
 632                                  penalty)