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 16import re 17 18from yapf_third_party._ylib2to3 import pytree 19from yapf_third_party._ylib2to3.pgen2 import token as grammar_token 20 21from yapf.pytree import pytree_utils 22from yapf.pytree import pytree_visitor 23from yapf.yapflib import style 24from yapf.yapflib import subtypes 25 26# TODO(morbo): Document the annotations in a centralized place. E.g., the 27# README file. 28UNBREAKABLE = 1000 * 1000 29NAMED_ASSIGN = 15000 30DOTTED_NAME = 4000 31VERY_STRONGLY_CONNECTED = 3500 32STRONGLY_CONNECTED = 3000 33CONNECTED = 500 34TOGETHER = 100 35 36OR_TEST = 1000 37AND_TEST = 1100 38NOT_TEST = 1200 39COMPARISON = 1300 40STAR_EXPR = 1300 41EXPR = 1400 42XOR_EXPR = 1500 43AND_EXPR = 1700 44SHIFT_EXPR = 1800 45ARITH_EXPR = 1900 46TERM = 2000 47FACTOR = 2100 48POWER = 2200 49ATOM = 2300 50ONE_ELEMENT_ARGUMENT = 500 51SUBSCRIPT = 6000 52 53 54def 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 63class _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 508def _SetUnbreakable(node): 509 """Set an UNBREAKABLE penalty annotation for the given node.""" 510 _RecAnnotate(node, pytree_utils.Annotation.SPLIT_PENALTY, UNBREAKABLE) 511 512 513def _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 520def _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 541def _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 553def _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 565def _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 585def _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 608def _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 623def _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 630def _SetSplitPenalty(node, penalty): 631 pytree_utils.SetNodeAnnotation(node, pytree_utils.Annotation.SPLIT_PENALTY, 632 penalty) 633