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