]> git.madduck.net Git - etc/vim.git/blob - src/black/linegen.py

madduck's git repository

Every one of the projects in this repository is available at the canonical URL git://git.madduck.net/madduck/pub/<projectpath> — see each project's metadata for the exact URL.

All patches and comments are welcome. Please squash your changes to logical commits before using git-format-patch and git-send-email to patches@git.madduck.net. If you'd read over the Git project's submission guidelines and adhered to them, I'd be especially grateful.

SSH access, as well as push access can be individually arranged.

If you use my repositories frequently, consider adding the following snippet to ~/.gitconfig and using the third clone URL listed for each project:

[url "git://git.madduck.net/madduck/"]
  insteadOf = madduck:

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