src/black/linegen.py

   1 """
   2 Generating lines of code.
   3 """
   4 from functools import partial, wraps
   5 import sys
   6 from typing import Collection, Iterator, List, Optional, Set, Union
   7
   8 from black.nodes import WHITESPACE, RARROW, STATEMENT, STANDALONE_COMMENT
   9 from black.nodes import ASSIGNMENTS, OPENING_BRACKETS, CLOSING_BRACKETS
  10 from black.nodes import Visitor, syms, is_arith_like, ensure_visible
  11 from black.nodes import (
  12     is_docstring,
  13     is_empty_tuple,
  14     is_one_tuple,
  15     is_one_sequence_between,
  16 )
  17 from black.nodes import is_name_token, is_lpar_token, is_rpar_token
  18 from black.nodes import is_walrus_assignment, is_yield, is_vararg, is_multiline_string
  19 from black.nodes import is_stub_suite, is_stub_body, is_atom_with_invisible_parens
  20 from black.nodes import wrap_in_parentheses
  21 from black.brackets import max_delimiter_priority_in_atom
  22 from black.brackets import DOT_PRIORITY, COMMA_PRIORITY
  23 from black.lines import Line, line_to_string, is_line_short_enough
  24 from black.lines import can_omit_invisible_parens, can_be_split, append_leaves
  25 from black.comments import generate_comments, list_comments, FMT_OFF
  26 from black.numerics import normalize_numeric_literal
  27 from black.strings import get_string_prefix, fix_docstring
  28 from black.strings import normalize_string_prefix, normalize_string_quotes
  29 from black.trans import Transformer, CannotTransform, StringMerger, StringSplitter
  30 from black.trans import StringParenWrapper, StringParenStripper, hug_power_op
  31 from black.mode import Mode, Feature, Preview
  32
  33 from blib2to3.pytree import Node, Leaf
  34 from blib2to3.pgen2 import token
  35
  36
  37 # types
  38 LeafID = int
  39 LN = Union[Leaf, Node]
  40
  41
  42 class CannotSplit(CannotTransform):
  43     """A readable split that fits the allotted line length is impossible."""
  44
  45
  46 # This isn't a dataclass because @dataclass + Generic breaks mypyc.
  47 # See also https://github.com/mypyc/mypyc/issues/827.
  48 class LineGenerator(Visitor[Line]):
  49     """Generates reformatted Line objects.  Empty lines are not emitted.
  50
  51     Note: destroys the tree it's visiting by mutating prefixes of its leaves
  52     in ways that will no longer stringify to valid Python code on the tree.
  53     """
  54
  55     def __init__(self, mode: Mode) -> None:
  56         self.mode = mode
  57         self.current_line: Line
  58         self.__post_init__()
  59
  60     def line(self, indent: int = 0) -> Iterator[Line]:
  61         """Generate a line.
  62
  63         If the line is empty, only emit if it makes sense.
  64         If the line is too long, split it first and then generate.
  65
  66         If any lines were generated, set up a new current_line.
  67         """
  68         if not self.current_line:
  69             self.current_line.depth += indent
  70             return  # Line is empty, don't emit. Creating a new one unnecessary.
  71
  72         complete_line = self.current_line
  73         self.current_line = Line(mode=self.mode, depth=complete_line.depth + indent)
  74         yield complete_line
  75
  76     def visit_default(self, node: LN) -> Iterator[Line]:
  77         """Default `visit_*()` implementation. Recurses to children of `node`."""
  78         if isinstance(node, Leaf):
  79             any_open_brackets = self.current_line.bracket_tracker.any_open_brackets()
  80             for comment in generate_comments(node, preview=self.mode.preview):
  81                 if any_open_brackets:
  82                     # any comment within brackets is subject to splitting
  83                     self.current_line.append(comment)
  84                 elif comment.type == token.COMMENT:
  85                     # regular trailing comment
  86                     self.current_line.append(comment)
  87                     yield from self.line()
  88
  89                 else:
  90                     # regular standalone comment
  91                     yield from self.line()
  92
  93                     self.current_line.append(comment)
  94                     yield from self.line()
  95
  96             normalize_prefix(node, inside_brackets=any_open_brackets)
  97             if self.mode.string_normalization and node.type == token.STRING:
  98                 node.value = normalize_string_prefix(node.value)
  99                 node.value = normalize_string_quotes(node.value)
 100             if node.type == token.NUMBER:
 101                 normalize_numeric_literal(node)
 102             if node.type not in WHITESPACE:
 103                 self.current_line.append(node)
 104         yield from super().visit_default(node)
 105
 106     def visit_INDENT(self, node: Leaf) -> Iterator[Line]:
 107         """Increase indentation level, maybe yield a line."""
 108         # In blib2to3 INDENT never holds comments.
 109         yield from self.line(+1)
 110         yield from self.visit_default(node)
 111
 112     def visit_DEDENT(self, node: Leaf) -> Iterator[Line]:
 113         """Decrease indentation level, maybe yield a line."""
 114         # The current line might still wait for trailing comments.  At DEDENT time
 115         # there won't be any (they would be prefixes on the preceding NEWLINE).
 116         # Emit the line then.
 117         yield from self.line()
 118
 119         # While DEDENT has no value, its prefix may contain standalone comments
 120         # that belong to the current indentation level.  Get 'em.
 121         yield from self.visit_default(node)
 122
 123         # Finally, emit the dedent.
 124         yield from self.line(-1)
 125
 126     def visit_stmt(
 127         self, node: Node, keywords: Set[str], parens: Set[str]
 128     ) -> Iterator[Line]:
 129         """Visit a statement.
 130
 131         This implementation is shared for `if`, `while`, `for`, `try`, `except`,
 132         `def`, `with`, `class`, `assert`, and assignments.
 133
 134         The relevant Python language `keywords` for a given statement will be
 135         NAME leaves within it. This methods puts those on a separate line.
 136
 137         `parens` holds a set of string leaf values immediately after which
 138         invisible parens should be put.
 139         """
 140         normalize_invisible_parens(node, parens_after=parens, preview=self.mode.preview)
 141         for child in node.children:
 142             if is_name_token(child) and child.value in keywords:
 143                 yield from self.line()
 144
 145             yield from self.visit(child)
 146
 147     def visit_match_case(self, node: Node) -> Iterator[Line]:
 148         """Visit either a match or case statement."""
 149         normalize_invisible_parens(node, parens_after=set(), preview=self.mode.preview)
 150
 151         yield from self.line()
 152         for child in node.children:
 153             yield from self.visit(child)
 154
 155     def visit_suite(self, node: Node) -> Iterator[Line]:
 156         """Visit a suite."""
 157         if self.mode.is_pyi and is_stub_suite(node):
 158             yield from self.visit(node.children[2])
 159         else:
 160             yield from self.visit_default(node)
 161
 162     def visit_simple_stmt(self, node: Node) -> Iterator[Line]:
 163         """Visit a statement without nested statements."""
 164         prev_type: Optional[int] = None
 165         for child in node.children:
 166             if (prev_type is None or prev_type == token.SEMI) and is_arith_like(child):
 167                 wrap_in_parentheses(node, child, visible=False)
 168             prev_type = child.type
 169
 170         is_suite_like = node.parent and node.parent.type in STATEMENT
 171         if is_suite_like:
 172             if self.mode.is_pyi and is_stub_body(node):
 173                 yield from self.visit_default(node)
 174             else:
 175                 yield from self.line(+1)
 176                 yield from self.visit_default(node)
 177                 yield from self.line(-1)
 178
 179         else:
 180             if (
 181                 not self.mode.is_pyi
 182                 or not node.parent
 183                 or not is_stub_suite(node.parent)
 184             ):
 185                 yield from self.line()
 186             yield from self.visit_default(node)
 187
 188     def visit_async_stmt(self, node: Node) -> Iterator[Line]:
 189         """Visit `async def`, `async for`, `async with`."""
 190         yield from self.line()
 191
 192         children = iter(node.children)
 193         for child in children:
 194             yield from self.visit(child)
 195
 196             if child.type == token.ASYNC:
 197                 break
 198
 199         internal_stmt = next(children)
 200         for child in internal_stmt.children:
 201             yield from self.visit(child)
 202
 203     def visit_decorators(self, node: Node) -> Iterator[Line]:
 204         """Visit decorators."""
 205         for child in node.children:
 206             yield from self.line()
 207             yield from self.visit(child)
 208
 209     def visit_power(self, node: Node) -> Iterator[Line]:
 210         for idx, leaf in enumerate(node.children[:-1]):
 211             next_leaf = node.children[idx + 1]
 212
 213             if not isinstance(leaf, Leaf):
 214                 continue
 215
 216             value = leaf.value.lower()
 217             if (
 218                 leaf.type == token.NUMBER
 219                 and next_leaf.type == syms.trailer
 220                 # Ensure that we are in an attribute trailer
 221                 and next_leaf.children[0].type == token.DOT
 222                 # It shouldn't wrap hexadecimal, binary and octal literals
 223                 and not value.startswith(("0x", "0b", "0o"))
 224                 # It shouldn't wrap complex literals
 225                 and "j" not in value
 226             ):
 227                 wrap_in_parentheses(node, leaf)
 228
 229         yield from self.visit_default(node)
 230
 231     def visit_SEMI(self, leaf: Leaf) -> Iterator[Line]:
 232         """Remove a semicolon and put the other statement on a separate line."""
 233         yield from self.line()
 234
 235     def visit_ENDMARKER(self, leaf: Leaf) -> Iterator[Line]:
 236         """End of file. Process outstanding comments and end with a newline."""
 237         yield from self.visit_default(leaf)
 238         yield from self.line()
 239
 240     def visit_STANDALONE_COMMENT(self, leaf: Leaf) -> Iterator[Line]:
 241         if not self.current_line.bracket_tracker.any_open_brackets():
 242             yield from self.line()
 243         yield from self.visit_default(leaf)
 244
 245     def visit_factor(self, node: Node) -> Iterator[Line]:
 246         """Force parentheses between a unary op and a binary power:
 247
 248         -2 ** 8 -> -(2 ** 8)
 249         """
 250         _operator, operand = node.children
 251         if (
 252             operand.type == syms.power
 253             and len(operand.children) == 3
 254             and operand.children[1].type == token.DOUBLESTAR
 255         ):
 256             lpar = Leaf(token.LPAR, "(")
 257             rpar = Leaf(token.RPAR, ")")
 258             index = operand.remove() or 0
 259             node.insert_child(index, Node(syms.atom, [lpar, operand, rpar]))
 260         yield from self.visit_default(node)
 261
 262     def visit_STRING(self, leaf: Leaf) -> Iterator[Line]:
 263         if is_docstring(leaf) and "\\\n" not in leaf.value:
 264             # We're ignoring docstrings with backslash newline escapes because changing
 265             # indentation of those changes the AST representation of the code.
 266             docstring = normalize_string_prefix(leaf.value)
 267             prefix = get_string_prefix(docstring)
 268             docstring = docstring[len(prefix) :]  # Remove the prefix
 269             quote_char = docstring[0]
 270             # A natural way to remove the outer quotes is to do:
 271             #   docstring = docstring.strip(quote_char)
 272             # but that breaks on """""x""" (which is '""x').
 273             # So we actually need to remove the first character and the next two
 274             # characters but only if they are the same as the first.
 275             quote_len = 1 if docstring[1] != quote_char else 3
 276             docstring = docstring[quote_len:-quote_len]
 277             docstring_started_empty = not docstring
 278
 279             if is_multiline_string(leaf):
 280                 indent = " " * 4 * self.current_line.depth
 281                 docstring = fix_docstring(docstring, indent)
 282             else:
 283                 docstring = docstring.strip()
 284
 285             if docstring:
 286                 # Add some padding if the docstring starts / ends with a quote mark.
 287                 if docstring[0] == quote_char:
 288                     docstring = " " + docstring
 289                 if docstring[-1] == quote_char:
 290                     docstring += " "
 291                 if docstring[-1] == "\\":
 292                     backslash_count = len(docstring) - len(docstring.rstrip("\\"))
 293                     if backslash_count % 2:
 294                         # Odd number of tailing backslashes, add some padding to
 295                         # avoid escaping the closing string quote.
 296                         docstring += " "
 297             elif not docstring_started_empty:
 298                 docstring = " "
 299
 300             # We could enforce triple quotes at this point.
 301             quote = quote_char * quote_len
 302             leaf.value = prefix + quote + docstring + quote
 303
 304         yield from self.visit_default(leaf)
 305
 306     def __post_init__(self) -> None:
 307         """You are in a twisty little maze of passages."""
 308         self.current_line = Line(mode=self.mode)
 309
 310         v = self.visit_stmt
 311         Ø: Set[str] = set()
 312         self.visit_assert_stmt = partial(v, keywords={"assert"}, parens={"assert", ","})
 313         self.visit_if_stmt = partial(
 314             v, keywords={"if", "else", "elif"}, parens={"if", "elif"}
 315         )
 316         self.visit_while_stmt = partial(v, keywords={"while", "else"}, parens={"while"})
 317         self.visit_for_stmt = partial(v, keywords={"for", "else"}, parens={"for", "in"})
 318         self.visit_try_stmt = partial(
 319             v, keywords={"try", "except", "else", "finally"}, parens=Ø
 320         )
 321         if self.mode.preview:
 322             self.visit_except_clause = partial(
 323                 v, keywords={"except"}, parens={"except"}
 324             )
 325             self.visit_with_stmt = partial(v, keywords={"with"}, parens={"with"})
 326         else:
 327             self.visit_except_clause = partial(v, keywords={"except"}, parens=Ø)
 328             self.visit_with_stmt = partial(v, keywords={"with"}, parens=Ø)
 329         self.visit_funcdef = partial(v, keywords={"def"}, parens=Ø)
 330         self.visit_classdef = partial(v, keywords={"class"}, parens=Ø)
 331         self.visit_expr_stmt = partial(v, keywords=Ø, parens=ASSIGNMENTS)
 332         self.visit_return_stmt = partial(v, keywords={"return"}, parens={"return"})
 333         self.visit_import_from = partial(v, keywords=Ø, parens={"import"})
 334         self.visit_del_stmt = partial(v, keywords=Ø, parens={"del"})
 335         self.visit_async_funcdef = self.visit_async_stmt
 336         self.visit_decorated = self.visit_decorators
 337
 338         # PEP 634
 339         self.visit_match_stmt = self.visit_match_case
 340         self.visit_case_block = self.visit_match_case
 341
 342
 343 def transform_line(
 344     line: Line, mode: Mode, features: Collection[Feature] = ()
 345 ) -> Iterator[Line]:
 346     """Transform a `line`, potentially splitting it into many lines.
 347
 348     They should fit in the allotted `line_length` but might not be able to.
 349
 350     `features` are syntactical features that may be used in the output.
 351     """
 352     if line.is_comment:
 353         yield line
 354         return
 355
 356     line_str = line_to_string(line)
 357
 358     ll = mode.line_length
 359     sn = mode.string_normalization
 360     string_merge = StringMerger(ll, sn)
 361     string_paren_strip = StringParenStripper(ll, sn)
 362     string_split = StringSplitter(ll, sn)
 363     string_paren_wrap = StringParenWrapper(ll, sn)
 364
 365     transformers: List[Transformer]
 366     if (
 367         not line.contains_uncollapsable_type_comments()
 368         and not line.should_split_rhs
 369         and not line.magic_trailing_comma
 370         and (
 371             is_line_short_enough(line, line_length=mode.line_length, line_str=line_str)
 372             or line.contains_unsplittable_type_ignore()
 373         )
 374         and not (line.inside_brackets and line.contains_standalone_comments())
 375     ):
 376         # Only apply basic string preprocessing, since lines shouldn't be split here.
 377         if Preview.string_processing in mode:
 378             transformers = [string_merge, string_paren_strip]
 379         else:
 380             transformers = []
 381     elif line.is_def:
 382         transformers = [left_hand_split]
 383     else:
 384
 385         def _rhs(
 386             self: object, line: Line, features: Collection[Feature]
 387         ) -> Iterator[Line]:
 388             """Wraps calls to `right_hand_split`.
 389
 390             The calls increasingly `omit` right-hand trailers (bracket pairs with
 391             content), meaning the trailers get glued together to split on another
 392             bracket pair instead.
 393             """
 394             for omit in generate_trailers_to_omit(line, mode.line_length):
 395                 lines = list(
 396                     right_hand_split(line, mode.line_length, features, omit=omit)
 397                 )
 398                 # Note: this check is only able to figure out if the first line of the
 399                 # *current* transformation fits in the line length.  This is true only
 400                 # for simple cases.  All others require running more transforms via
 401                 # `transform_line()`.  This check doesn't know if those would succeed.
 402                 if is_line_short_enough(lines[0], line_length=mode.line_length):
 403                     yield from lines
 404                     return
 405
 406             # All splits failed, best effort split with no omits.
 407             # This mostly happens to multiline strings that are by definition
 408             # reported as not fitting a single line, as well as lines that contain
 409             # trailing commas (those have to be exploded).
 410             yield from right_hand_split(
 411                 line, line_length=mode.line_length, features=features
 412             )
 413
 414         # HACK: nested functions (like _rhs) compiled by mypyc don't retain their
 415         # __name__ attribute which is needed in `run_transformer` further down.
 416         # Unfortunately a nested class breaks mypyc too. So a class must be created
 417         # via type ... https://github.com/mypyc/mypyc/issues/884
 418         rhs = type("rhs", (), {"__call__": _rhs})()
 419
 420         if Preview.string_processing in mode:
 421             if line.inside_brackets:
 422                 transformers = [
 423                     string_merge,
 424                     string_paren_strip,
 425                     string_split,
 426                     delimiter_split,
 427                     standalone_comment_split,
 428                     string_paren_wrap,
 429                     rhs,
 430                 ]
 431             else:
 432                 transformers = [
 433                     string_merge,
 434                     string_paren_strip,
 435                     string_split,
 436                     string_paren_wrap,
 437                     rhs,
 438                 ]
 439         else:
 440             if line.inside_brackets:
 441                 transformers = [delimiter_split, standalone_comment_split, rhs]
 442             else:
 443                 transformers = [rhs]
 444     # It's always safe to attempt hugging of power operations and pretty much every line
 445     # could match.
 446     transformers.append(hug_power_op)
 447
 448     for transform in transformers:
 449         # We are accumulating lines in `result` because we might want to abort
 450         # mission and return the original line in the end, or attempt a different
 451         # split altogether.
 452         try:
 453             result = run_transformer(line, transform, mode, features, line_str=line_str)
 454         except CannotTransform:
 455             continue
 456         else:
 457             yield from result
 458             break
 459
 460     else:
 461         yield line
 462
 463
 464 def left_hand_split(line: Line, _features: Collection[Feature] = ()) -> Iterator[Line]:
 465     """Split line into many lines, starting with the first matching bracket pair.
 466
 467     Note: this usually looks weird, only use this for function definitions.
 468     Prefer RHS otherwise.  This is why this function is not symmetrical with
 469     :func:`right_hand_split` which also handles optional parentheses.
 470     """
 471     tail_leaves: List[Leaf] = []
 472     body_leaves: List[Leaf] = []
 473     head_leaves: List[Leaf] = []
 474     current_leaves = head_leaves
 475     matching_bracket: Optional[Leaf] = None
 476     for leaf in line.leaves:
 477         if (
 478             current_leaves is body_leaves
 479             and leaf.type in CLOSING_BRACKETS
 480             and leaf.opening_bracket is matching_bracket
 481         ):
 482             current_leaves = tail_leaves if body_leaves else head_leaves
 483         current_leaves.append(leaf)
 484         if current_leaves is head_leaves:
 485             if leaf.type in OPENING_BRACKETS:
 486                 matching_bracket = leaf
 487                 current_leaves = body_leaves
 488     if not matching_bracket:
 489         raise CannotSplit("No brackets found")
 490
 491     head = bracket_split_build_line(head_leaves, line, matching_bracket)
 492     body = bracket_split_build_line(body_leaves, line, matching_bracket, is_body=True)
 493     tail = bracket_split_build_line(tail_leaves, line, matching_bracket)
 494     bracket_split_succeeded_or_raise(head, body, tail)
 495     for result in (head, body, tail):
 496         if result:
 497             yield result
 498
 499
 500 def right_hand_split(
 501     line: Line,
 502     line_length: int,
 503     features: Collection[Feature] = (),
 504     omit: Collection[LeafID] = (),
 505 ) -> Iterator[Line]:
 506     """Split line into many lines, starting with the last matching bracket pair.
 507
 508     If the split was by optional parentheses, attempt splitting without them, too.
 509     `omit` is a collection of closing bracket IDs that shouldn't be considered for
 510     this split.
 511
 512     Note: running this function modifies `bracket_depth` on the leaves of `line`.
 513     """
 514     tail_leaves: List[Leaf] = []
 515     body_leaves: List[Leaf] = []
 516     head_leaves: List[Leaf] = []
 517     current_leaves = tail_leaves
 518     opening_bracket: Optional[Leaf] = None
 519     closing_bracket: Optional[Leaf] = None
 520     for leaf in reversed(line.leaves):
 521         if current_leaves is body_leaves:
 522             if leaf is opening_bracket:
 523                 current_leaves = head_leaves if body_leaves else tail_leaves
 524         current_leaves.append(leaf)
 525         if current_leaves is tail_leaves:
 526             if leaf.type in CLOSING_BRACKETS and id(leaf) not in omit:
 527                 opening_bracket = leaf.opening_bracket
 528                 closing_bracket = leaf
 529                 current_leaves = body_leaves
 530     if not (opening_bracket and closing_bracket and head_leaves):
 531         # If there is no opening or closing_bracket that means the split failed and
 532         # all content is in the tail.  Otherwise, if `head_leaves` are empty, it means
 533         # the matching `opening_bracket` wasn't available on `line` anymore.
 534         raise CannotSplit("No brackets found")
 535
 536     tail_leaves.reverse()
 537     body_leaves.reverse()
 538     head_leaves.reverse()
 539     head = bracket_split_build_line(head_leaves, line, opening_bracket)
 540     body = bracket_split_build_line(body_leaves, line, opening_bracket, is_body=True)
 541     tail = bracket_split_build_line(tail_leaves, line, opening_bracket)
 542     bracket_split_succeeded_or_raise(head, body, tail)
 543     if (
 544         Feature.FORCE_OPTIONAL_PARENTHESES not in features
 545         # the opening bracket is an optional paren
 546         and opening_bracket.type == token.LPAR
 547         and not opening_bracket.value
 548         # the closing bracket is an optional paren
 549         and closing_bracket.type == token.RPAR
 550         and not closing_bracket.value
 551         # it's not an import (optional parens are the only thing we can split on
 552         # in this case; attempting a split without them is a waste of time)
 553         and not line.is_import
 554         # there are no standalone comments in the body
 555         and not body.contains_standalone_comments(0)
 556         # and we can actually remove the parens
 557         and can_omit_invisible_parens(body, line_length)
 558     ):
 559         omit = {id(closing_bracket), *omit}
 560         try:
 561             yield from right_hand_split(line, line_length, features=features, omit=omit)
 562             return
 563
 564         except CannotSplit as e:
 565             if not (
 566                 can_be_split(body)
 567                 or is_line_short_enough(body, line_length=line_length)
 568             ):
 569                 raise CannotSplit(
 570                     "Splitting failed, body is still too long and can't be split."
 571                 ) from e
 572
 573             elif head.contains_multiline_strings() or tail.contains_multiline_strings():
 574                 raise CannotSplit(
 575                     "The current optional pair of parentheses is bound to fail to"
 576                     " satisfy the splitting algorithm because the head or the tail"
 577                     " contains multiline strings which by definition never fit one"
 578                     " line."
 579                 ) from e
 580
 581     ensure_visible(opening_bracket)
 582     ensure_visible(closing_bracket)
 583     for result in (head, body, tail):
 584         if result:
 585             yield result
 586
 587
 588 def bracket_split_succeeded_or_raise(head: Line, body: Line, tail: Line) -> None:
 589     """Raise :exc:`CannotSplit` if the last left- or right-hand split failed.
 590
 591     Do nothing otherwise.
 592
 593     A left- or right-hand split is based on a pair of brackets. Content before
 594     (and including) the opening bracket is left on one line, content inside the
 595     brackets is put on a separate line, and finally content starting with and
 596     following the closing bracket is put on a separate line.
 597
 598     Those are called `head`, `body`, and `tail`, respectively. If the split
 599     produced the same line (all content in `head`) or ended up with an empty `body`
 600     and the `tail` is just the closing bracket, then it's considered failed.
 601     """
 602     tail_len = len(str(tail).strip())
 603     if not body:
 604         if tail_len == 0:
 605             raise CannotSplit("Splitting brackets produced the same line")
 606
 607         elif tail_len < 3:
 608             raise CannotSplit(
 609                 f"Splitting brackets on an empty body to save {tail_len} characters is"
 610                 " not worth it"
 611             )
 612
 613
 614 def bracket_split_build_line(
 615     leaves: List[Leaf], original: Line, opening_bracket: Leaf, *, is_body: bool = False
 616 ) -> Line:
 617     """Return a new line with given `leaves` and respective comments from `original`.
 618
 619     If `is_body` is True, the result line is one-indented inside brackets and as such
 620     has its first leaf's prefix normalized and a trailing comma added when expected.
 621     """
 622     result = Line(mode=original.mode, depth=original.depth)
 623     if is_body:
 624         result.inside_brackets = True
 625         result.depth += 1
 626         if leaves:
 627             # Since body is a new indent level, remove spurious leading whitespace.
 628             normalize_prefix(leaves[0], inside_brackets=True)
 629             # Ensure a trailing comma for imports and standalone function arguments, but
 630             # be careful not to add one after any comments or within type annotations.
 631             no_commas = (
 632                 original.is_def
 633                 and opening_bracket.value == "("
 634                 and not any(leaf.type == token.COMMA for leaf in leaves)
 635                 # In particular, don't add one within a parenthesized return annotation.
 636                 # Unfortunately the indicator we're in a return annotation (RARROW) may
 637                 # be defined directly in the parent node, the parent of the parent ...
 638                 # and so on depending on how complex the return annotation is.
 639                 # This isn't perfect and there's some false negatives but they are in
 640                 # contexts were a comma is actually fine.
 641                 and not any(
 642                     node.prev_sibling.type == RARROW
 643                     for node in (
 644                         leaves[0].parent,
 645                         getattr(leaves[0].parent, "parent", None),
 646                     )
 647                     if isinstance(node, Node) and isinstance(node.prev_sibling, Leaf)
 648                 )
 649             )
 650
 651             if original.is_import or no_commas:
 652                 for i in range(len(leaves) - 1, -1, -1):
 653                     if leaves[i].type == STANDALONE_COMMENT:
 654                         continue
 655
 656                     if leaves[i].type != token.COMMA:
 657                         new_comma = Leaf(token.COMMA, ",")
 658                         leaves.insert(i + 1, new_comma)
 659                     break
 660
 661     # Populate the line
 662     for leaf in leaves:
 663         result.append(leaf, preformatted=True)
 664         for comment_after in original.comments_after(leaf):
 665             result.append(comment_after, preformatted=True)
 666     if is_body and should_split_line(result, opening_bracket):
 667         result.should_split_rhs = True
 668     return result
 669
 670
 671 def dont_increase_indentation(split_func: Transformer) -> Transformer:
 672     """Normalize prefix of the first leaf in every line returned by `split_func`.
 673
 674     This is a decorator over relevant split functions.
 675     """
 676
 677     @wraps(split_func)
 678     def split_wrapper(line: Line, features: Collection[Feature] = ()) -> Iterator[Line]:
 679         for split_line in split_func(line, features):
 680             normalize_prefix(split_line.leaves[0], inside_brackets=True)
 681             yield split_line
 682
 683     return split_wrapper
 684
 685
 686 @dont_increase_indentation
 687 def delimiter_split(line: Line, features: Collection[Feature] = ()) -> Iterator[Line]:
 688     """Split according to delimiters of the highest priority.
 689
 690     If the appropriate Features are given, the split will add trailing commas
 691     also in function signatures and calls that contain `*` and `**`.
 692     """
 693     try:
 694         last_leaf = line.leaves[-1]
 695     except IndexError:
 696         raise CannotSplit("Line empty") from None
 697
 698     bt = line.bracket_tracker
 699     try:
 700         delimiter_priority = bt.max_delimiter_priority(exclude={id(last_leaf)})
 701     except ValueError:
 702         raise CannotSplit("No delimiters found") from None
 703
 704     if delimiter_priority == DOT_PRIORITY:
 705         if bt.delimiter_count_with_priority(delimiter_priority) == 1:
 706             raise CannotSplit("Splitting a single attribute from its owner looks wrong")
 707
 708     current_line = Line(
 709         mode=line.mode, depth=line.depth, inside_brackets=line.inside_brackets
 710     )
 711     lowest_depth = sys.maxsize
 712     trailing_comma_safe = True
 713
 714     def append_to_line(leaf: Leaf) -> Iterator[Line]:
 715         """Append `leaf` to current line or to new line if appending impossible."""
 716         nonlocal current_line
 717         try:
 718             current_line.append_safe(leaf, preformatted=True)
 719         except ValueError:
 720             yield current_line
 721
 722             current_line = Line(
 723                 mode=line.mode, depth=line.depth, inside_brackets=line.inside_brackets
 724             )
 725             current_line.append(leaf)
 726
 727     for leaf in line.leaves:
 728         yield from append_to_line(leaf)
 729
 730         for comment_after in line.comments_after(leaf):
 731             yield from append_to_line(comment_after)
 732
 733         lowest_depth = min(lowest_depth, leaf.bracket_depth)
 734         if leaf.bracket_depth == lowest_depth:
 735             if is_vararg(leaf, within={syms.typedargslist}):
 736                 trailing_comma_safe = (
 737                     trailing_comma_safe and Feature.TRAILING_COMMA_IN_DEF in features
 738                 )
 739             elif is_vararg(leaf, within={syms.arglist, syms.argument}):
 740                 trailing_comma_safe = (
 741                     trailing_comma_safe and Feature.TRAILING_COMMA_IN_CALL in features
 742                 )
 743
 744         leaf_priority = bt.delimiters.get(id(leaf))
 745         if leaf_priority == delimiter_priority:
 746             yield current_line
 747
 748             current_line = Line(
 749                 mode=line.mode, depth=line.depth, inside_brackets=line.inside_brackets
 750             )
 751     if current_line:
 752         if (
 753             trailing_comma_safe
 754             and delimiter_priority == COMMA_PRIORITY
 755             and current_line.leaves[-1].type != token.COMMA
 756             and current_line.leaves[-1].type != STANDALONE_COMMENT
 757         ):
 758             new_comma = Leaf(token.COMMA, ",")
 759             current_line.append(new_comma)
 760         yield current_line
 761
 762
 763 @dont_increase_indentation
 764 def standalone_comment_split(
 765     line: Line, features: Collection[Feature] = ()
 766 ) -> Iterator[Line]:
 767     """Split standalone comments from the rest of the line."""
 768     if not line.contains_standalone_comments(0):
 769         raise CannotSplit("Line does not have any standalone comments")
 770
 771     current_line = Line(
 772         mode=line.mode, depth=line.depth, inside_brackets=line.inside_brackets
 773     )
 774
 775     def append_to_line(leaf: Leaf) -> Iterator[Line]:
 776         """Append `leaf` to current line or to new line if appending impossible."""
 777         nonlocal current_line
 778         try:
 779             current_line.append_safe(leaf, preformatted=True)
 780         except ValueError:
 781             yield current_line
 782
 783             current_line = Line(
 784                 line.mode, depth=line.depth, inside_brackets=line.inside_brackets
 785             )
 786             current_line.append(leaf)
 787
 788     for leaf in line.leaves:
 789         yield from append_to_line(leaf)
 790
 791         for comment_after in line.comments_after(leaf):
 792             yield from append_to_line(comment_after)
 793
 794     if current_line:
 795         yield current_line
 796
 797
 798 def normalize_prefix(leaf: Leaf, *, inside_brackets: bool) -> None:
 799     """Leave existing extra newlines if not `inside_brackets`. Remove everything
 800     else.
 801
 802     Note: don't use backslashes for formatting or you'll lose your voting rights.
 803     """
 804     if not inside_brackets:
 805         spl = leaf.prefix.split("#")
 806         if "\\" not in spl[0]:
 807             nl_count = spl[-1].count("\n")
 808             if len(spl) > 1:
 809                 nl_count -= 1
 810             leaf.prefix = "\n" * nl_count
 811             return
 812
 813     leaf.prefix = ""
 814
 815
 816 def normalize_invisible_parens(
 817     node: Node, parens_after: Set[str], *, preview: bool
 818 ) -> None:
 819     """Make existing optional parentheses invisible or create new ones.
 820
 821     `parens_after` is a set of string leaf values immediately after which parens
 822     should be put.
 823
 824     Standardizes on visible parentheses for single-element tuples, and keeps
 825     existing visible parentheses for other tuples and generator expressions.
 826     """
 827     for pc in list_comments(node.prefix, is_endmarker=False, preview=preview):
 828         if pc.value in FMT_OFF:
 829             # This `node` has a prefix with `# fmt: off`, don't mess with parens.
 830             return
 831     check_lpar = False
 832     for index, child in enumerate(list(node.children)):
 833         # Fixes a bug where invisible parens are not properly stripped from
 834         # assignment statements that contain type annotations.
 835         if isinstance(child, Node) and child.type == syms.annassign:
 836             normalize_invisible_parens(
 837                 child, parens_after=parens_after, preview=preview
 838             )
 839
 840         # Add parentheses around long tuple unpacking in assignments.
 841         if (
 842             index == 0
 843             and isinstance(child, Node)
 844             and child.type == syms.testlist_star_expr
 845         ):
 846             check_lpar = True
 847
 848         if check_lpar:
 849             if (
 850                 preview
 851                 and child.type == syms.atom
 852                 and node.type == syms.for_stmt
 853                 and isinstance(child.prev_sibling, Leaf)
 854                 and child.prev_sibling.type == token.NAME
 855                 and child.prev_sibling.value == "for"
 856             ):
 857                 if maybe_make_parens_invisible_in_atom(
 858                     child,
 859                     parent=node,
 860                     remove_brackets_around_comma=True,
 861                 ):
 862                     wrap_in_parentheses(node, child, visible=False)
 863             elif preview and isinstance(child, Node) and node.type == syms.with_stmt:
 864                 remove_with_parens(child, node)
 865             elif child.type == syms.atom:
 866                 if maybe_make_parens_invisible_in_atom(
 867                     child,
 868                     parent=node,
 869                 ):
 870                     wrap_in_parentheses(node, child, visible=False)
 871             elif is_one_tuple(child):
 872                 wrap_in_parentheses(node, child, visible=True)
 873             elif node.type == syms.import_from:
 874                 # "import from" nodes store parentheses directly as part of
 875                 # the statement
 876                 if is_lpar_token(child):
 877                     assert is_rpar_token(node.children[-1])
 878                     # make parentheses invisible
 879                     child.value = ""
 880                     node.children[-1].value = ""
 881                 elif child.type != token.STAR:
 882                     # insert invisible parentheses
 883                     node.insert_child(index, Leaf(token.LPAR, ""))
 884                     node.append_child(Leaf(token.RPAR, ""))
 885                 break
 886
 887             elif not (isinstance(child, Leaf) and is_multiline_string(child)):
 888                 wrap_in_parentheses(node, child, visible=False)
 889
 890         comma_check = child.type == token.COMMA if preview else False
 891
 892         check_lpar = isinstance(child, Leaf) and (
 893             child.value in parens_after or comma_check
 894         )
 895
 896
 897 def remove_with_parens(node: Node, parent: Node) -> None:
 898     """Recursively hide optional parens in `with` statements."""
 899     # Removing all unnecessary parentheses in with statements in one pass is a tad
 900     # complex as different variations of bracketed statements result in pretty
 901     # different parse trees:
 902     #
 903     # with (open("file")) as f:                       # this is an asexpr_test
 904     #     ...
 905     #
 906     # with (open("file") as f):                       # this is an atom containing an
 907     #     ...                                         # asexpr_test
 908     #
 909     # with (open("file")) as f, (open("file")) as f:  # this is asexpr_test, COMMA,
 910     #     ...                                         # asexpr_test
 911     #
 912     # with (open("file") as f, open("file") as f):    # an atom containing a
 913     #     ...                                         # testlist_gexp which then
 914     #                                                 # contains multiple asexpr_test(s)
 915     if node.type == syms.atom:
 916         if maybe_make_parens_invisible_in_atom(
 917             node,
 918             parent=parent,
 919             remove_brackets_around_comma=True,
 920         ):
 921             wrap_in_parentheses(parent, node, visible=False)
 922         if isinstance(node.children[1], Node):
 923             remove_with_parens(node.children[1], node)
 924     elif node.type == syms.testlist_gexp:
 925         for child in node.children:
 926             if isinstance(child, Node):
 927                 remove_with_parens(child, node)
 928     elif node.type == syms.asexpr_test and not any(
 929         leaf.type == token.COLONEQUAL for leaf in node.leaves()
 930     ):
 931         if maybe_make_parens_invisible_in_atom(
 932             node.children[0],
 933             parent=node,
 934             remove_brackets_around_comma=True,
 935         ):
 936             wrap_in_parentheses(node, node.children[0], visible=False)
 937
 938
 939 def maybe_make_parens_invisible_in_atom(
 940     node: LN,
 941     parent: LN,
 942     remove_brackets_around_comma: bool = False,
 943 ) -> bool:
 944     """If it's safe, make the parens in the atom `node` invisible, recursively.
 945     Additionally, remove repeated, adjacent invisible parens from the atom `node`
 946     as they are redundant.
 947
 948     Returns whether the node should itself be wrapped in invisible parentheses.
 949     """
 950     if (
 951         node.type != syms.atom
 952         or is_empty_tuple(node)
 953         or is_one_tuple(node)
 954         or (is_yield(node) and parent.type != syms.expr_stmt)
 955         or (
 956             # This condition tries to prevent removing non-optional brackets
 957             # around a tuple, however, can be a bit overzealous so we provide
 958             # and option to skip this check for `for` and `with` statements.
 959             not remove_brackets_around_comma
 960             and max_delimiter_priority_in_atom(node) >= COMMA_PRIORITY
 961         )
 962     ):
 963         return False
 964
 965     if is_walrus_assignment(node):
 966         if parent.type in [
 967             syms.annassign,
 968             syms.expr_stmt,
 969             syms.assert_stmt,
 970             syms.return_stmt,
 971             # these ones aren't useful to end users, but they do please fuzzers
 972             syms.for_stmt,
 973             syms.del_stmt,
 974         ]:
 975             return False
 976
 977     first = node.children[0]
 978     last = node.children[-1]
 979     if is_lpar_token(first) and is_rpar_token(last):
 980         middle = node.children[1]
 981         # make parentheses invisible
 982         first.value = ""
 983         last.value = ""
 984         maybe_make_parens_invisible_in_atom(
 985             middle,
 986             parent=parent,
 987             remove_brackets_around_comma=remove_brackets_around_comma,
 988         )
 989
 990         if is_atom_with_invisible_parens(middle):
 991             # Strip the invisible parens from `middle` by replacing
 992             # it with the child in-between the invisible parens
 993             middle.replace(middle.children[1])
 994
 995         return False
 996
 997     return True
 998
 999
1000 def should_split_line(line: Line, opening_bracket: Leaf) -> bool:
1001     """Should `line` be immediately split with `delimiter_split()` after RHS?"""
1002
1003     if not (opening_bracket.parent and opening_bracket.value in "[{("):
1004         return False
1005
1006     # We're essentially checking if the body is delimited by commas and there's more
1007     # than one of them (we're excluding the trailing comma and if the delimiter priority
1008     # is still commas, that means there's more).
1009     exclude = set()
1010     trailing_comma = False
1011     try:
1012         last_leaf = line.leaves[-1]
1013         if last_leaf.type == token.COMMA:
1014             trailing_comma = True
1015             exclude.add(id(last_leaf))
1016         max_priority = line.bracket_tracker.max_delimiter_priority(exclude=exclude)
1017     except (IndexError, ValueError):
1018         return False
1019
1020     return max_priority == COMMA_PRIORITY and (
1021         (line.mode.magic_trailing_comma and trailing_comma)
1022         # always explode imports
1023         or opening_bracket.parent.type in {syms.atom, syms.import_from}
1024     )
1025
1026
1027 def generate_trailers_to_omit(line: Line, line_length: int) -> Iterator[Set[LeafID]]:
1028     """Generate sets of closing bracket IDs that should be omitted in a RHS.
1029
1030     Brackets can be omitted if the entire trailer up to and including
1031     a preceding closing bracket fits in one line.
1032
1033     Yielded sets are cumulative (contain results of previous yields, too).  First
1034     set is empty, unless the line should explode, in which case bracket pairs until
1035     the one that needs to explode are omitted.
1036     """
1037
1038     omit: Set[LeafID] = set()
1039     if not line.magic_trailing_comma:
1040         yield omit
1041
1042     length = 4 * line.depth
1043     opening_bracket: Optional[Leaf] = None
1044     closing_bracket: Optional[Leaf] = None
1045     inner_brackets: Set[LeafID] = set()
1046     for index, leaf, leaf_length in line.enumerate_with_length(reversed=True):
1047         length += leaf_length
1048         if length > line_length:
1049             break
1050
1051         has_inline_comment = leaf_length > len(leaf.value) + len(leaf.prefix)
1052         if leaf.type == STANDALONE_COMMENT or has_inline_comment:
1053             break
1054
1055         if opening_bracket:
1056             if leaf is opening_bracket:
1057                 opening_bracket = None
1058             elif leaf.type in CLOSING_BRACKETS:
1059                 prev = line.leaves[index - 1] if index > 0 else None
1060                 if (
1061                     prev
1062                     and prev.type == token.COMMA
1063                     and leaf.opening_bracket is not None
1064                     and not is_one_sequence_between(
1065                         leaf.opening_bracket, leaf, line.leaves
1066                     )
1067                 ):
1068                     # Never omit bracket pairs with trailing commas.
1069                     # We need to explode on those.
1070                     break
1071
1072                 inner_brackets.add(id(leaf))
1073         elif leaf.type in CLOSING_BRACKETS:
1074             prev = line.leaves[index - 1] if index > 0 else None
1075             if prev and prev.type in OPENING_BRACKETS:
1076                 # Empty brackets would fail a split so treat them as "inner"
1077                 # brackets (e.g. only add them to the `omit` set if another
1078                 # pair of brackets was good enough.
1079                 inner_brackets.add(id(leaf))
1080                 continue
1081
1082             if closing_bracket:
1083                 omit.add(id(closing_bracket))
1084                 omit.update(inner_brackets)
1085                 inner_brackets.clear()
1086                 yield omit
1087
1088             if (
1089                 prev
1090                 and prev.type == token.COMMA
1091                 and leaf.opening_bracket is not None
1092                 and not is_one_sequence_between(leaf.opening_bracket, leaf, line.leaves)
1093             ):
1094                 # Never omit bracket pairs with trailing commas.
1095                 # We need to explode on those.
1096                 break
1097
1098             if leaf.value:
1099                 opening_bracket = leaf.opening_bracket
1100                 closing_bracket = leaf
1101
1102
1103 def run_transformer(
1104     line: Line,
1105     transform: Transformer,
1106     mode: Mode,
1107     features: Collection[Feature],
1108     *,
1109     line_str: str = "",
1110 ) -> List[Line]:
1111     if not line_str:
1112         line_str = line_to_string(line)
1113     result: List[Line] = []
1114     for transformed_line in transform(line, features):
1115         if str(transformed_line).strip("\n") == line_str:
1116             raise CannotTransform("Line transformer returned an unchanged result")
1117
1118         result.extend(transform_line(transformed_line, mode=mode, features=features))
1119
1120     if (
1121         transform.__class__.__name__ != "rhs"
1122         or not line.bracket_tracker.invisible
1123         or any(bracket.value for bracket in line.bracket_tracker.invisible)
1124         or line.contains_multiline_strings()
1125         or result[0].contains_uncollapsable_type_comments()
1126         or result[0].contains_unsplittable_type_ignore()
1127         or is_line_short_enough(result[0], line_length=mode.line_length)
1128         # If any leaves have no parents (which _can_ occur since
1129         # `transform(line)` potentially destroys the line's underlying node
1130         # structure), then we can't proceed. Doing so would cause the below
1131         # call to `append_leaves()` to fail.
1132         or any(leaf.parent is None for leaf in line.leaves)
1133     ):
1134         return result
1135
1136     line_copy = line.clone()
1137     append_leaves(line_copy, line, line.leaves)
1138     features_fop = set(features) | {Feature.FORCE_OPTIONAL_PARENTHESES}
1139     second_opinion = run_transformer(
1140         line_copy, transform, mode, features_fop, line_str=line_str
1141     )
1142     if all(
1143         is_line_short_enough(ln, line_length=mode.line_length) for ln in second_opinion
1144     ):
1145         result = second_opinion
1146     return result