]> 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:

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