]> git.madduck.net Git - etc/vim.git/blob - black.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:

0aba2c5de0fd6f20b48f17aa057f0bbac3c0d22f
[etc/vim.git] / black.py
1 import asyncio
2 import pickle
3 from asyncio.base_events import BaseEventLoop
4 from concurrent.futures import Executor, ProcessPoolExecutor
5 from enum import Enum
6 from functools import partial, wraps
7 import keyword
8 import logging
9 from multiprocessing import Manager
10 import os
11 from pathlib import Path
12 import re
13 import tokenize
14 import signal
15 import sys
16 from typing import (
17     Any,
18     Callable,
19     Collection,
20     Dict,
21     Generic,
22     Iterable,
23     Iterator,
24     List,
25     Optional,
26     Pattern,
27     Sequence,
28     Set,
29     Tuple,
30     Type,
31     TypeVar,
32     Union,
33     cast,
34 )
35
36 from appdirs import user_cache_dir
37 from attr import dataclass, Factory
38 import click
39
40 # lib2to3 fork
41 from blib2to3.pytree import Node, Leaf, type_repr
42 from blib2to3 import pygram, pytree
43 from blib2to3.pgen2 import driver, token
44 from blib2to3.pgen2.parse import ParseError
45
46
47 __version__ = "18.5b0"
48 DEFAULT_LINE_LENGTH = 88
49 CACHE_DIR = Path(user_cache_dir("black", version=__version__))
50
51
52 # types
53 FileContent = str
54 Encoding = str
55 Depth = int
56 NodeType = int
57 LeafID = int
58 Priority = int
59 Index = int
60 LN = Union[Leaf, Node]
61 SplitFunc = Callable[["Line", bool], Iterator["Line"]]
62 Timestamp = float
63 FileSize = int
64 CacheInfo = Tuple[Timestamp, FileSize]
65 Cache = Dict[Path, CacheInfo]
66 out = partial(click.secho, bold=True, err=True)
67 err = partial(click.secho, fg="red", err=True)
68
69 pygram.initialize(CACHE_DIR)
70 syms = pygram.python_symbols
71
72
73 class NothingChanged(UserWarning):
74     """Raised by :func:`format_file` when reformatted code is the same as source."""
75
76
77 class CannotSplit(Exception):
78     """A readable split that fits the allotted line length is impossible.
79
80     Raised by :func:`left_hand_split`, :func:`right_hand_split`, and
81     :func:`delimiter_split`.
82     """
83
84
85 class FormatError(Exception):
86     """Base exception for `# fmt: on` and `# fmt: off` handling.
87
88     It holds the number of bytes of the prefix consumed before the format
89     control comment appeared.
90     """
91
92     def __init__(self, consumed: int) -> None:
93         super().__init__(consumed)
94         self.consumed = consumed
95
96     def trim_prefix(self, leaf: Leaf) -> None:
97         leaf.prefix = leaf.prefix[self.consumed :]
98
99     def leaf_from_consumed(self, leaf: Leaf) -> Leaf:
100         """Returns a new Leaf from the consumed part of the prefix."""
101         unformatted_prefix = leaf.prefix[: self.consumed]
102         return Leaf(token.NEWLINE, unformatted_prefix)
103
104
105 class FormatOn(FormatError):
106     """Found a comment like `# fmt: on` in the file."""
107
108
109 class FormatOff(FormatError):
110     """Found a comment like `# fmt: off` in the file."""
111
112
113 class WriteBack(Enum):
114     NO = 0
115     YES = 1
116     DIFF = 2
117
118
119 class Changed(Enum):
120     NO = 0
121     CACHED = 1
122     YES = 2
123
124
125 @click.command()
126 @click.option(
127     "-l",
128     "--line-length",
129     type=int,
130     default=DEFAULT_LINE_LENGTH,
131     help="How many character per line to allow.",
132     show_default=True,
133 )
134 @click.option(
135     "--check",
136     is_flag=True,
137     help=(
138         "Don't write the files back, just return the status.  Return code 0 "
139         "means nothing would change.  Return code 1 means some files would be "
140         "reformatted.  Return code 123 means there was an internal error."
141     ),
142 )
143 @click.option(
144     "--diff",
145     is_flag=True,
146     help="Don't write the files back, just output a diff for each file on stdout.",
147 )
148 @click.option(
149     "--fast/--safe",
150     is_flag=True,
151     help="If --fast given, skip temporary sanity checks. [default: --safe]",
152 )
153 @click.option(
154     "-q",
155     "--quiet",
156     is_flag=True,
157     help=(
158         "Don't emit non-error messages to stderr. Errors are still emitted, "
159         "silence those with 2>/dev/null."
160     ),
161 )
162 @click.option(
163     "--pyi",
164     is_flag=True,
165     help="Force pyi (stub file) formatting, regardless of file extension.",
166 )
167 @click.option(
168     "--py36",
169     is_flag=True,
170     help=(
171         "Force Python 3.6 mode, even if file doesn't currently use "
172         "Python 3.6-only syntax."
173     ),
174 )
175 @click.version_option(version=__version__)
176 @click.argument(
177     "src",
178     nargs=-1,
179     type=click.Path(
180         exists=True, file_okay=True, dir_okay=True, readable=True, allow_dash=True
181     ),
182 )
183 @click.pass_context
184 def main(
185     ctx: click.Context,
186     line_length: int,
187     check: bool,
188     diff: bool,
189     fast: bool,
190     pyi: bool,
191     py36: bool,
192     quiet: bool,
193     src: List[str],
194 ) -> None:
195     """The uncompromising code formatter."""
196     sources: List[Path] = []
197     for s in src:
198         p = Path(s)
199         if p.is_dir():
200             sources.extend(gen_python_files_in_dir(p))
201         elif p.is_file():
202             # if a file was explicitly given, we don't care about its extension
203             sources.append(p)
204         elif s == "-":
205             sources.append(Path("-"))
206         else:
207             err(f"invalid path: {s}")
208
209     if check and not diff:
210         write_back = WriteBack.NO
211     elif diff:
212         write_back = WriteBack.DIFF
213     else:
214         write_back = WriteBack.YES
215     report = Report(check=check, quiet=quiet)
216     if len(sources) == 0:
217         out("No paths given. Nothing to do 😴")
218         ctx.exit(0)
219         return
220
221     elif len(sources) == 1:
222         reformat_one(
223             src=sources[0],
224             line_length=line_length,
225             fast=fast,
226             pyi=pyi,
227             py36=py36,
228             write_back=write_back,
229             report=report,
230         )
231     else:
232         loop = asyncio.get_event_loop()
233         executor = ProcessPoolExecutor(max_workers=os.cpu_count())
234         try:
235             loop.run_until_complete(
236                 schedule_formatting(
237                     sources=sources,
238                     line_length=line_length,
239                     fast=fast,
240                     pyi=pyi,
241                     py36=py36,
242                     write_back=write_back,
243                     report=report,
244                     loop=loop,
245                     executor=executor,
246                 )
247             )
248         finally:
249             shutdown(loop)
250         if not quiet:
251             out("All done! ✨ 🍰 ✨")
252             click.echo(str(report))
253     ctx.exit(report.return_code)
254
255
256 def reformat_one(
257     src: Path,
258     line_length: int,
259     fast: bool,
260     pyi: bool,
261     py36: bool,
262     write_back: WriteBack,
263     report: "Report",
264 ) -> None:
265     """Reformat a single file under `src` without spawning child processes.
266
267     If `quiet` is True, non-error messages are not output. `line_length`,
268     `write_back`, `fast` and `pyi` options are passed to
269     :func:`format_file_in_place` or :func:`format_stdin_to_stdout`.
270     """
271     try:
272         changed = Changed.NO
273         if not src.is_file() and str(src) == "-":
274             if format_stdin_to_stdout(
275                 line_length=line_length,
276                 fast=fast,
277                 is_pyi=pyi,
278                 force_py36=py36,
279                 write_back=write_back,
280             ):
281                 changed = Changed.YES
282         else:
283             cache: Cache = {}
284             if write_back != WriteBack.DIFF:
285                 cache = read_cache(line_length, pyi, py36)
286                 src = src.resolve()
287                 if src in cache and cache[src] == get_cache_info(src):
288                     changed = Changed.CACHED
289             if changed is not Changed.CACHED and format_file_in_place(
290                 src,
291                 line_length=line_length,
292                 fast=fast,
293                 force_pyi=pyi,
294                 force_py36=py36,
295                 write_back=write_back,
296             ):
297                 changed = Changed.YES
298             if write_back == WriteBack.YES and changed is not Changed.NO:
299                 write_cache(cache, [src], line_length, pyi, py36)
300         report.done(src, changed)
301     except Exception as exc:
302         report.failed(src, str(exc))
303
304
305 async def schedule_formatting(
306     sources: List[Path],
307     line_length: int,
308     fast: bool,
309     pyi: bool,
310     py36: bool,
311     write_back: WriteBack,
312     report: "Report",
313     loop: BaseEventLoop,
314     executor: Executor,
315 ) -> None:
316     """Run formatting of `sources` in parallel using the provided `executor`.
317
318     (Use ProcessPoolExecutors for actual parallelism.)
319
320     `line_length`, `write_back`, `fast`, and `pyi` options are passed to
321     :func:`format_file_in_place`.
322     """
323     cache: Cache = {}
324     if write_back != WriteBack.DIFF:
325         cache = read_cache(line_length, pyi, py36)
326         sources, cached = filter_cached(cache, sources)
327         for src in cached:
328             report.done(src, Changed.CACHED)
329     cancelled = []
330     formatted = []
331     if sources:
332         lock = None
333         if write_back == WriteBack.DIFF:
334             # For diff output, we need locks to ensure we don't interleave output
335             # from different processes.
336             manager = Manager()
337             lock = manager.Lock()
338         tasks = {
339             loop.run_in_executor(
340                 executor,
341                 format_file_in_place,
342                 src,
343                 line_length,
344                 fast,
345                 pyi,
346                 py36,
347                 write_back,
348                 lock,
349             ): src
350             for src in sorted(sources)
351         }
352         pending: Iterable[asyncio.Task] = tasks.keys()
353         try:
354             loop.add_signal_handler(signal.SIGINT, cancel, pending)
355             loop.add_signal_handler(signal.SIGTERM, cancel, pending)
356         except NotImplementedError:
357             # There are no good alternatives for these on Windows
358             pass
359         while pending:
360             done, _ = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
361             for task in done:
362                 src = tasks.pop(task)
363                 if task.cancelled():
364                     cancelled.append(task)
365                 elif task.exception():
366                     report.failed(src, str(task.exception()))
367                 else:
368                     formatted.append(src)
369                     report.done(src, Changed.YES if task.result() else Changed.NO)
370     if cancelled:
371         await asyncio.gather(*cancelled, loop=loop, return_exceptions=True)
372     if write_back == WriteBack.YES and formatted:
373         write_cache(cache, formatted, line_length, pyi, py36)
374
375
376 def format_file_in_place(
377     src: Path,
378     line_length: int,
379     fast: bool,
380     force_pyi: bool = False,
381     force_py36: bool = False,
382     write_back: WriteBack = WriteBack.NO,
383     lock: Any = None,  # multiprocessing.Manager().Lock() is some crazy proxy
384 ) -> bool:
385     """Format file under `src` path. Return True if changed.
386
387     If `write_back` is True, write reformatted code back to stdout.
388     `line_length` and `fast` options are passed to :func:`format_file_contents`.
389     """
390     is_pyi = force_pyi or src.suffix == ".pyi"
391
392     with tokenize.open(src) as src_buffer:
393         src_contents = src_buffer.read()
394     try:
395         dst_contents = format_file_contents(
396             src_contents,
397             line_length=line_length,
398             fast=fast,
399             is_pyi=is_pyi,
400             force_py36=force_py36,
401         )
402     except NothingChanged:
403         return False
404
405     if write_back == write_back.YES:
406         with open(src, "w", encoding=src_buffer.encoding) as f:
407             f.write(dst_contents)
408     elif write_back == write_back.DIFF:
409         src_name = f"{src}  (original)"
410         dst_name = f"{src}  (formatted)"
411         diff_contents = diff(src_contents, dst_contents, src_name, dst_name)
412         if lock:
413             lock.acquire()
414         try:
415             sys.stdout.write(diff_contents)
416         finally:
417             if lock:
418                 lock.release()
419     return True
420
421
422 def format_stdin_to_stdout(
423     line_length: int,
424     fast: bool,
425     is_pyi: bool = False,
426     force_py36: bool = False,
427     write_back: WriteBack = WriteBack.NO,
428 ) -> bool:
429     """Format file on stdin. Return True if changed.
430
431     If `write_back` is True, write reformatted code back to stdout.
432     `line_length`, `fast`, `is_pyi`, and `force_py36` arguments are passed to
433     :func:`format_file_contents`.
434     """
435     src = sys.stdin.read()
436     dst = src
437     try:
438         dst = format_file_contents(
439             src,
440             line_length=line_length,
441             fast=fast,
442             is_pyi=is_pyi,
443             force_py36=force_py36,
444         )
445         return True
446
447     except NothingChanged:
448         return False
449
450     finally:
451         if write_back == WriteBack.YES:
452             sys.stdout.write(dst)
453         elif write_back == WriteBack.DIFF:
454             src_name = "<stdin>  (original)"
455             dst_name = "<stdin>  (formatted)"
456             sys.stdout.write(diff(src, dst, src_name, dst_name))
457
458
459 def format_file_contents(
460     src_contents: str,
461     *,
462     line_length: int,
463     fast: bool,
464     is_pyi: bool = False,
465     force_py36: bool = False,
466 ) -> FileContent:
467     """Reformat contents a file and return new contents.
468
469     If `fast` is False, additionally confirm that the reformatted code is
470     valid by calling :func:`assert_equivalent` and :func:`assert_stable` on it.
471     `line_length` is passed to :func:`format_str`.
472     """
473     if src_contents.strip() == "":
474         raise NothingChanged
475
476     dst_contents = format_str(
477         src_contents, line_length=line_length, is_pyi=is_pyi, force_py36=force_py36
478     )
479     if src_contents == dst_contents:
480         raise NothingChanged
481
482     if not fast:
483         assert_equivalent(src_contents, dst_contents)
484         assert_stable(
485             src_contents,
486             dst_contents,
487             line_length=line_length,
488             is_pyi=is_pyi,
489             force_py36=force_py36,
490         )
491     return dst_contents
492
493
494 def format_str(
495     src_contents: str,
496     line_length: int,
497     *,
498     is_pyi: bool = False,
499     force_py36: bool = False,
500 ) -> FileContent:
501     """Reformat a string and return new contents.
502
503     `line_length` determines how many characters per line are allowed.
504     """
505     src_node = lib2to3_parse(src_contents)
506     dst_contents = ""
507     future_imports = get_future_imports(src_node)
508     elt = EmptyLineTracker(is_pyi=is_pyi)
509     py36 = force_py36 or is_python36(src_node)
510     lines = LineGenerator(
511         remove_u_prefix=py36 or "unicode_literals" in future_imports, is_pyi=is_pyi
512     )
513     empty_line = Line()
514     after = 0
515     for current_line in lines.visit(src_node):
516         for _ in range(after):
517             dst_contents += str(empty_line)
518         before, after = elt.maybe_empty_lines(current_line)
519         for _ in range(before):
520             dst_contents += str(empty_line)
521         for line in split_line(current_line, line_length=line_length, py36=py36):
522             dst_contents += str(line)
523     return dst_contents
524
525
526 GRAMMARS = [
527     pygram.python_grammar_no_print_statement_no_exec_statement,
528     pygram.python_grammar_no_print_statement,
529     pygram.python_grammar,
530 ]
531
532
533 def lib2to3_parse(src_txt: str) -> Node:
534     """Given a string with source, return the lib2to3 Node."""
535     grammar = pygram.python_grammar_no_print_statement
536     if src_txt[-1] != "\n":
537         nl = "\r\n" if "\r\n" in src_txt[:1024] else "\n"
538         src_txt += nl
539     for grammar in GRAMMARS:
540         drv = driver.Driver(grammar, pytree.convert)
541         try:
542             result = drv.parse_string(src_txt, True)
543             break
544
545         except ParseError as pe:
546             lineno, column = pe.context[1]
547             lines = src_txt.splitlines()
548             try:
549                 faulty_line = lines[lineno - 1]
550             except IndexError:
551                 faulty_line = "<line number missing in source>"
552             exc = ValueError(f"Cannot parse: {lineno}:{column}: {faulty_line}")
553     else:
554         raise exc from None
555
556     if isinstance(result, Leaf):
557         result = Node(syms.file_input, [result])
558     return result
559
560
561 def lib2to3_unparse(node: Node) -> str:
562     """Given a lib2to3 node, return its string representation."""
563     code = str(node)
564     return code
565
566
567 T = TypeVar("T")
568
569
570 class Visitor(Generic[T]):
571     """Basic lib2to3 visitor that yields things of type `T` on `visit()`."""
572
573     def visit(self, node: LN) -> Iterator[T]:
574         """Main method to visit `node` and its children.
575
576         It tries to find a `visit_*()` method for the given `node.type`, like
577         `visit_simple_stmt` for Node objects or `visit_INDENT` for Leaf objects.
578         If no dedicated `visit_*()` method is found, chooses `visit_default()`
579         instead.
580
581         Then yields objects of type `T` from the selected visitor.
582         """
583         if node.type < 256:
584             name = token.tok_name[node.type]
585         else:
586             name = type_repr(node.type)
587         yield from getattr(self, f"visit_{name}", self.visit_default)(node)
588
589     def visit_default(self, node: LN) -> Iterator[T]:
590         """Default `visit_*()` implementation. Recurses to children of `node`."""
591         if isinstance(node, Node):
592             for child in node.children:
593                 yield from self.visit(child)
594
595
596 @dataclass
597 class DebugVisitor(Visitor[T]):
598     tree_depth: int = 0
599
600     def visit_default(self, node: LN) -> Iterator[T]:
601         indent = " " * (2 * self.tree_depth)
602         if isinstance(node, Node):
603             _type = type_repr(node.type)
604             out(f"{indent}{_type}", fg="yellow")
605             self.tree_depth += 1
606             for child in node.children:
607                 yield from self.visit(child)
608
609             self.tree_depth -= 1
610             out(f"{indent}/{_type}", fg="yellow", bold=False)
611         else:
612             _type = token.tok_name.get(node.type, str(node.type))
613             out(f"{indent}{_type}", fg="blue", nl=False)
614             if node.prefix:
615                 # We don't have to handle prefixes for `Node` objects since
616                 # that delegates to the first child anyway.
617                 out(f" {node.prefix!r}", fg="green", bold=False, nl=False)
618             out(f" {node.value!r}", fg="blue", bold=False)
619
620     @classmethod
621     def show(cls, code: str) -> None:
622         """Pretty-print the lib2to3 AST of a given string of `code`.
623
624         Convenience method for debugging.
625         """
626         v: DebugVisitor[None] = DebugVisitor()
627         list(v.visit(lib2to3_parse(code)))
628
629
630 KEYWORDS = set(keyword.kwlist)
631 WHITESPACE = {token.DEDENT, token.INDENT, token.NEWLINE}
632 FLOW_CONTROL = {"return", "raise", "break", "continue"}
633 STATEMENT = {
634     syms.if_stmt,
635     syms.while_stmt,
636     syms.for_stmt,
637     syms.try_stmt,
638     syms.except_clause,
639     syms.with_stmt,
640     syms.funcdef,
641     syms.classdef,
642 }
643 STANDALONE_COMMENT = 153
644 LOGIC_OPERATORS = {"and", "or"}
645 COMPARATORS = {
646     token.LESS,
647     token.GREATER,
648     token.EQEQUAL,
649     token.NOTEQUAL,
650     token.LESSEQUAL,
651     token.GREATEREQUAL,
652 }
653 MATH_OPERATORS = {
654     token.VBAR,
655     token.CIRCUMFLEX,
656     token.AMPER,
657     token.LEFTSHIFT,
658     token.RIGHTSHIFT,
659     token.PLUS,
660     token.MINUS,
661     token.STAR,
662     token.SLASH,
663     token.DOUBLESLASH,
664     token.PERCENT,
665     token.AT,
666     token.TILDE,
667     token.DOUBLESTAR,
668 }
669 STARS = {token.STAR, token.DOUBLESTAR}
670 VARARGS_PARENTS = {
671     syms.arglist,
672     syms.argument,  # double star in arglist
673     syms.trailer,  # single argument to call
674     syms.typedargslist,
675     syms.varargslist,  # lambdas
676 }
677 UNPACKING_PARENTS = {
678     syms.atom,  # single element of a list or set literal
679     syms.dictsetmaker,
680     syms.listmaker,
681     syms.testlist_gexp,
682 }
683 TEST_DESCENDANTS = {
684     syms.test,
685     syms.lambdef,
686     syms.or_test,
687     syms.and_test,
688     syms.not_test,
689     syms.comparison,
690     syms.star_expr,
691     syms.expr,
692     syms.xor_expr,
693     syms.and_expr,
694     syms.shift_expr,
695     syms.arith_expr,
696     syms.trailer,
697     syms.term,
698     syms.power,
699 }
700 ASSIGNMENTS = {
701     "=",
702     "+=",
703     "-=",
704     "*=",
705     "@=",
706     "/=",
707     "%=",
708     "&=",
709     "|=",
710     "^=",
711     "<<=",
712     ">>=",
713     "**=",
714     "//=",
715 }
716 COMPREHENSION_PRIORITY = 20
717 COMMA_PRIORITY = 18
718 TERNARY_PRIORITY = 16
719 LOGIC_PRIORITY = 14
720 STRING_PRIORITY = 12
721 COMPARATOR_PRIORITY = 10
722 MATH_PRIORITIES = {
723     token.VBAR: 9,
724     token.CIRCUMFLEX: 8,
725     token.AMPER: 7,
726     token.LEFTSHIFT: 6,
727     token.RIGHTSHIFT: 6,
728     token.PLUS: 5,
729     token.MINUS: 5,
730     token.STAR: 4,
731     token.SLASH: 4,
732     token.DOUBLESLASH: 4,
733     token.PERCENT: 4,
734     token.AT: 4,
735     token.TILDE: 3,
736     token.DOUBLESTAR: 2,
737 }
738 DOT_PRIORITY = 1
739
740
741 @dataclass
742 class BracketTracker:
743     """Keeps track of brackets on a line."""
744
745     depth: int = 0
746     bracket_match: Dict[Tuple[Depth, NodeType], Leaf] = Factory(dict)
747     delimiters: Dict[LeafID, Priority] = Factory(dict)
748     previous: Optional[Leaf] = None
749     _for_loop_variable: int = 0
750     _lambda_arguments: int = 0
751
752     def mark(self, leaf: Leaf) -> None:
753         """Mark `leaf` with bracket-related metadata. Keep track of delimiters.
754
755         All leaves receive an int `bracket_depth` field that stores how deep
756         within brackets a given leaf is. 0 means there are no enclosing brackets
757         that started on this line.
758
759         If a leaf is itself a closing bracket, it receives an `opening_bracket`
760         field that it forms a pair with. This is a one-directional link to
761         avoid reference cycles.
762
763         If a leaf is a delimiter (a token on which Black can split the line if
764         needed) and it's on depth 0, its `id()` is stored in the tracker's
765         `delimiters` field.
766         """
767         if leaf.type == token.COMMENT:
768             return
769
770         self.maybe_decrement_after_for_loop_variable(leaf)
771         self.maybe_decrement_after_lambda_arguments(leaf)
772         if leaf.type in CLOSING_BRACKETS:
773             self.depth -= 1
774             opening_bracket = self.bracket_match.pop((self.depth, leaf.type))
775             leaf.opening_bracket = opening_bracket
776         leaf.bracket_depth = self.depth
777         if self.depth == 0:
778             delim = is_split_before_delimiter(leaf, self.previous)
779             if delim and self.previous is not None:
780                 self.delimiters[id(self.previous)] = delim
781             else:
782                 delim = is_split_after_delimiter(leaf, self.previous)
783                 if delim:
784                     self.delimiters[id(leaf)] = delim
785         if leaf.type in OPENING_BRACKETS:
786             self.bracket_match[self.depth, BRACKET[leaf.type]] = leaf
787             self.depth += 1
788         self.previous = leaf
789         self.maybe_increment_lambda_arguments(leaf)
790         self.maybe_increment_for_loop_variable(leaf)
791
792     def any_open_brackets(self) -> bool:
793         """Return True if there is an yet unmatched open bracket on the line."""
794         return bool(self.bracket_match)
795
796     def max_delimiter_priority(self, exclude: Iterable[LeafID] = ()) -> int:
797         """Return the highest priority of a delimiter found on the line.
798
799         Values are consistent with what `is_split_*_delimiter()` return.
800         Raises ValueError on no delimiters.
801         """
802         return max(v for k, v in self.delimiters.items() if k not in exclude)
803
804     def delimiter_count_with_priority(self, priority: int = 0) -> int:
805         """Return the number of delimiters with the given `priority`.
806
807         If no `priority` is passed, defaults to max priority on the line.
808         """
809         if not self.delimiters:
810             return 0
811
812         priority = priority or self.max_delimiter_priority()
813         return sum(1 for p in self.delimiters.values() if p == priority)
814
815     def maybe_increment_for_loop_variable(self, leaf: Leaf) -> bool:
816         """In a for loop, or comprehension, the variables are often unpacks.
817
818         To avoid splitting on the comma in this situation, increase the depth of
819         tokens between `for` and `in`.
820         """
821         if leaf.type == token.NAME and leaf.value == "for":
822             self.depth += 1
823             self._for_loop_variable += 1
824             return True
825
826         return False
827
828     def maybe_decrement_after_for_loop_variable(self, leaf: Leaf) -> bool:
829         """See `maybe_increment_for_loop_variable` above for explanation."""
830         if self._for_loop_variable and leaf.type == token.NAME and leaf.value == "in":
831             self.depth -= 1
832             self._for_loop_variable -= 1
833             return True
834
835         return False
836
837     def maybe_increment_lambda_arguments(self, leaf: Leaf) -> bool:
838         """In a lambda expression, there might be more than one argument.
839
840         To avoid splitting on the comma in this situation, increase the depth of
841         tokens between `lambda` and `:`.
842         """
843         if leaf.type == token.NAME and leaf.value == "lambda":
844             self.depth += 1
845             self._lambda_arguments += 1
846             return True
847
848         return False
849
850     def maybe_decrement_after_lambda_arguments(self, leaf: Leaf) -> bool:
851         """See `maybe_increment_lambda_arguments` above for explanation."""
852         if self._lambda_arguments and leaf.type == token.COLON:
853             self.depth -= 1
854             self._lambda_arguments -= 1
855             return True
856
857         return False
858
859     def get_open_lsqb(self) -> Optional[Leaf]:
860         """Return the most recent opening square bracket (if any)."""
861         return self.bracket_match.get((self.depth - 1, token.RSQB))
862
863
864 @dataclass
865 class Line:
866     """Holds leaves and comments. Can be printed with `str(line)`."""
867
868     depth: int = 0
869     leaves: List[Leaf] = Factory(list)
870     comments: List[Tuple[Index, Leaf]] = Factory(list)
871     bracket_tracker: BracketTracker = Factory(BracketTracker)
872     inside_brackets: bool = False
873     should_explode: bool = False
874
875     def append(self, leaf: Leaf, preformatted: bool = False) -> None:
876         """Add a new `leaf` to the end of the line.
877
878         Unless `preformatted` is True, the `leaf` will receive a new consistent
879         whitespace prefix and metadata applied by :class:`BracketTracker`.
880         Trailing commas are maybe removed, unpacked for loop variables are
881         demoted from being delimiters.
882
883         Inline comments are put aside.
884         """
885         has_value = leaf.type in BRACKETS or bool(leaf.value.strip())
886         if not has_value:
887             return
888
889         if token.COLON == leaf.type and self.is_class_paren_empty:
890             del self.leaves[-2:]
891         if self.leaves and not preformatted:
892             # Note: at this point leaf.prefix should be empty except for
893             # imports, for which we only preserve newlines.
894             leaf.prefix += whitespace(
895                 leaf, complex_subscript=self.is_complex_subscript(leaf)
896             )
897         if self.inside_brackets or not preformatted:
898             self.bracket_tracker.mark(leaf)
899             self.maybe_remove_trailing_comma(leaf)
900         if not self.append_comment(leaf):
901             self.leaves.append(leaf)
902
903     def append_safe(self, leaf: Leaf, preformatted: bool = False) -> None:
904         """Like :func:`append()` but disallow invalid standalone comment structure.
905
906         Raises ValueError when any `leaf` is appended after a standalone comment
907         or when a standalone comment is not the first leaf on the line.
908         """
909         if self.bracket_tracker.depth == 0:
910             if self.is_comment:
911                 raise ValueError("cannot append to standalone comments")
912
913             if self.leaves and leaf.type == STANDALONE_COMMENT:
914                 raise ValueError(
915                     "cannot append standalone comments to a populated line"
916                 )
917
918         self.append(leaf, preformatted=preformatted)
919
920     @property
921     def is_comment(self) -> bool:
922         """Is this line a standalone comment?"""
923         return len(self.leaves) == 1 and self.leaves[0].type == STANDALONE_COMMENT
924
925     @property
926     def is_decorator(self) -> bool:
927         """Is this line a decorator?"""
928         return bool(self) and self.leaves[0].type == token.AT
929
930     @property
931     def is_import(self) -> bool:
932         """Is this an import line?"""
933         return bool(self) and is_import(self.leaves[0])
934
935     @property
936     def is_class(self) -> bool:
937         """Is this line a class definition?"""
938         return (
939             bool(self)
940             and self.leaves[0].type == token.NAME
941             and self.leaves[0].value == "class"
942         )
943
944     @property
945     def is_stub_class(self) -> bool:
946         """Is this line a class definition with a body consisting only of "..."?"""
947         return self.is_class and self.leaves[-3:] == [
948             Leaf(token.DOT, ".") for _ in range(3)
949         ]
950
951     @property
952     def is_def(self) -> bool:
953         """Is this a function definition? (Also returns True for async defs.)"""
954         try:
955             first_leaf = self.leaves[0]
956         except IndexError:
957             return False
958
959         try:
960             second_leaf: Optional[Leaf] = self.leaves[1]
961         except IndexError:
962             second_leaf = None
963         return (first_leaf.type == token.NAME and first_leaf.value == "def") or (
964             first_leaf.type == token.ASYNC
965             and second_leaf is not None
966             and second_leaf.type == token.NAME
967             and second_leaf.value == "def"
968         )
969
970     @property
971     def is_flow_control(self) -> bool:
972         """Is this line a flow control statement?
973
974         Those are `return`, `raise`, `break`, and `continue`.
975         """
976         return (
977             bool(self)
978             and self.leaves[0].type == token.NAME
979             and self.leaves[0].value in FLOW_CONTROL
980         )
981
982     @property
983     def is_yield(self) -> bool:
984         """Is this line a yield statement?"""
985         return (
986             bool(self)
987             and self.leaves[0].type == token.NAME
988             and self.leaves[0].value == "yield"
989         )
990
991     @property
992     def is_class_paren_empty(self) -> bool:
993         """Is this a class with no base classes but using parentheses?
994
995         Those are unnecessary and should be removed.
996         """
997         return (
998             bool(self)
999             and len(self.leaves) == 4
1000             and self.is_class
1001             and self.leaves[2].type == token.LPAR
1002             and self.leaves[2].value == "("
1003             and self.leaves[3].type == token.RPAR
1004             and self.leaves[3].value == ")"
1005         )
1006
1007     def contains_standalone_comments(self, depth_limit: int = sys.maxsize) -> bool:
1008         """If so, needs to be split before emitting."""
1009         for leaf in self.leaves:
1010             if leaf.type == STANDALONE_COMMENT:
1011                 if leaf.bracket_depth <= depth_limit:
1012                     return True
1013
1014         return False
1015
1016     def maybe_remove_trailing_comma(self, closing: Leaf) -> bool:
1017         """Remove trailing comma if there is one and it's safe."""
1018         if not (
1019             self.leaves
1020             and self.leaves[-1].type == token.COMMA
1021             and closing.type in CLOSING_BRACKETS
1022         ):
1023             return False
1024
1025         if closing.type == token.RBRACE:
1026             self.remove_trailing_comma()
1027             return True
1028
1029         if closing.type == token.RSQB:
1030             comma = self.leaves[-1]
1031             if comma.parent and comma.parent.type == syms.listmaker:
1032                 self.remove_trailing_comma()
1033                 return True
1034
1035         # For parens let's check if it's safe to remove the comma.
1036         # Imports are always safe.
1037         if self.is_import:
1038             self.remove_trailing_comma()
1039             return True
1040
1041         # Otheriwsse, if the trailing one is the only one, we might mistakenly
1042         # change a tuple into a different type by removing the comma.
1043         depth = closing.bracket_depth + 1
1044         commas = 0
1045         opening = closing.opening_bracket
1046         for _opening_index, leaf in enumerate(self.leaves):
1047             if leaf is opening:
1048                 break
1049
1050         else:
1051             return False
1052
1053         for leaf in self.leaves[_opening_index + 1 :]:
1054             if leaf is closing:
1055                 break
1056
1057             bracket_depth = leaf.bracket_depth
1058             if bracket_depth == depth and leaf.type == token.COMMA:
1059                 commas += 1
1060                 if leaf.parent and leaf.parent.type == syms.arglist:
1061                     commas += 1
1062                     break
1063
1064         if commas > 1:
1065             self.remove_trailing_comma()
1066             return True
1067
1068         return False
1069
1070     def append_comment(self, comment: Leaf) -> bool:
1071         """Add an inline or standalone comment to the line."""
1072         if (
1073             comment.type == STANDALONE_COMMENT
1074             and self.bracket_tracker.any_open_brackets()
1075         ):
1076             comment.prefix = ""
1077             return False
1078
1079         if comment.type != token.COMMENT:
1080             return False
1081
1082         after = len(self.leaves) - 1
1083         if after == -1:
1084             comment.type = STANDALONE_COMMENT
1085             comment.prefix = ""
1086             return False
1087
1088         else:
1089             self.comments.append((after, comment))
1090             return True
1091
1092     def comments_after(self, leaf: Leaf, _index: int = -1) -> Iterator[Leaf]:
1093         """Generate comments that should appear directly after `leaf`.
1094
1095         Provide a non-negative leaf `_index` to speed up the function.
1096         """
1097         if _index == -1:
1098             for _index, _leaf in enumerate(self.leaves):
1099                 if leaf is _leaf:
1100                     break
1101
1102             else:
1103                 return
1104
1105         for index, comment_after in self.comments:
1106             if _index == index:
1107                 yield comment_after
1108
1109     def remove_trailing_comma(self) -> None:
1110         """Remove the trailing comma and moves the comments attached to it."""
1111         comma_index = len(self.leaves) - 1
1112         for i in range(len(self.comments)):
1113             comment_index, comment = self.comments[i]
1114             if comment_index == comma_index:
1115                 self.comments[i] = (comma_index - 1, comment)
1116         self.leaves.pop()
1117
1118     def is_complex_subscript(self, leaf: Leaf) -> bool:
1119         """Return True iff `leaf` is part of a slice with non-trivial exprs."""
1120         open_lsqb = (
1121             leaf if leaf.type == token.LSQB else self.bracket_tracker.get_open_lsqb()
1122         )
1123         if open_lsqb is None:
1124             return False
1125
1126         subscript_start = open_lsqb.next_sibling
1127         if (
1128             isinstance(subscript_start, Node)
1129             and subscript_start.type == syms.subscriptlist
1130         ):
1131             subscript_start = child_towards(subscript_start, leaf)
1132         return subscript_start is not None and any(
1133             n.type in TEST_DESCENDANTS for n in subscript_start.pre_order()
1134         )
1135
1136     def __str__(self) -> str:
1137         """Render the line."""
1138         if not self:
1139             return "\n"
1140
1141         indent = "    " * self.depth
1142         leaves = iter(self.leaves)
1143         first = next(leaves)
1144         res = f"{first.prefix}{indent}{first.value}"
1145         for leaf in leaves:
1146             res += str(leaf)
1147         for _, comment in self.comments:
1148             res += str(comment)
1149         return res + "\n"
1150
1151     def __bool__(self) -> bool:
1152         """Return True if the line has leaves or comments."""
1153         return bool(self.leaves or self.comments)
1154
1155
1156 class UnformattedLines(Line):
1157     """Just like :class:`Line` but stores lines which aren't reformatted."""
1158
1159     def append(self, leaf: Leaf, preformatted: bool = True) -> None:
1160         """Just add a new `leaf` to the end of the lines.
1161
1162         The `preformatted` argument is ignored.
1163
1164         Keeps track of indentation `depth`, which is useful when the user
1165         says `# fmt: on`. Otherwise, doesn't do anything with the `leaf`.
1166         """
1167         try:
1168             list(generate_comments(leaf))
1169         except FormatOn as f_on:
1170             self.leaves.append(f_on.leaf_from_consumed(leaf))
1171             raise
1172
1173         self.leaves.append(leaf)
1174         if leaf.type == token.INDENT:
1175             self.depth += 1
1176         elif leaf.type == token.DEDENT:
1177             self.depth -= 1
1178
1179     def __str__(self) -> str:
1180         """Render unformatted lines from leaves which were added with `append()`.
1181
1182         `depth` is not used for indentation in this case.
1183         """
1184         if not self:
1185             return "\n"
1186
1187         res = ""
1188         for leaf in self.leaves:
1189             res += str(leaf)
1190         return res
1191
1192     def append_comment(self, comment: Leaf) -> bool:
1193         """Not implemented in this class. Raises `NotImplementedError`."""
1194         raise NotImplementedError("Unformatted lines don't store comments separately.")
1195
1196     def maybe_remove_trailing_comma(self, closing: Leaf) -> bool:
1197         """Does nothing and returns False."""
1198         return False
1199
1200     def maybe_increment_for_loop_variable(self, leaf: Leaf) -> bool:
1201         """Does nothing and returns False."""
1202         return False
1203
1204
1205 @dataclass
1206 class EmptyLineTracker:
1207     """Provides a stateful method that returns the number of potential extra
1208     empty lines needed before and after the currently processed line.
1209
1210     Note: this tracker works on lines that haven't been split yet.  It assumes
1211     the prefix of the first leaf consists of optional newlines.  Those newlines
1212     are consumed by `maybe_empty_lines()` and included in the computation.
1213     """
1214     is_pyi: bool = False
1215     previous_line: Optional[Line] = None
1216     previous_after: int = 0
1217     previous_defs: List[int] = Factory(list)
1218
1219     def maybe_empty_lines(self, current_line: Line) -> Tuple[int, int]:
1220         """Return the number of extra empty lines before and after the `current_line`.
1221
1222         This is for separating `def`, `async def` and `class` with extra empty
1223         lines (two on module-level), as well as providing an extra empty line
1224         after flow control keywords to make them more prominent.
1225         """
1226         if isinstance(current_line, UnformattedLines):
1227             return 0, 0
1228
1229         before, after = self._maybe_empty_lines(current_line)
1230         before -= self.previous_after
1231         self.previous_after = after
1232         self.previous_line = current_line
1233         return before, after
1234
1235     def _maybe_empty_lines(self, current_line: Line) -> Tuple[int, int]:
1236         max_allowed = 1
1237         if current_line.depth == 0:
1238             max_allowed = 1 if self.is_pyi else 2
1239         if current_line.leaves:
1240             # Consume the first leaf's extra newlines.
1241             first_leaf = current_line.leaves[0]
1242             before = first_leaf.prefix.count("\n")
1243             before = min(before, max_allowed)
1244             first_leaf.prefix = ""
1245         else:
1246             before = 0
1247         depth = current_line.depth
1248         while self.previous_defs and self.previous_defs[-1] >= depth:
1249             self.previous_defs.pop()
1250             if self.is_pyi:
1251                 before = 0 if depth else 1
1252             else:
1253                 before = 1 if depth else 2
1254         is_decorator = current_line.is_decorator
1255         if is_decorator or current_line.is_def or current_line.is_class:
1256             if not is_decorator:
1257                 self.previous_defs.append(depth)
1258             if self.previous_line is None:
1259                 # Don't insert empty lines before the first line in the file.
1260                 return 0, 0
1261
1262             if self.previous_line.is_decorator:
1263                 return 0, 0
1264
1265             if (
1266                 self.previous_line.is_comment
1267                 and self.previous_line.depth == current_line.depth
1268                 and before == 0
1269             ):
1270                 return 0, 0
1271
1272             if self.is_pyi:
1273                 if self.previous_line.depth > current_line.depth:
1274                     newlines = 1
1275                 elif current_line.is_class or self.previous_line.is_class:
1276                     if current_line.is_stub_class and self.previous_line.is_stub_class:
1277                         newlines = 0
1278                     else:
1279                         newlines = 1
1280                 else:
1281                     newlines = 0
1282             else:
1283                 newlines = 2
1284             if current_line.depth and newlines:
1285                 newlines -= 1
1286             return newlines, 0
1287
1288         if (
1289             self.previous_line
1290             and self.previous_line.is_import
1291             and not current_line.is_import
1292             and depth == self.previous_line.depth
1293         ):
1294             return (before or 1), 0
1295
1296         return before, 0
1297
1298
1299 @dataclass
1300 class LineGenerator(Visitor[Line]):
1301     """Generates reformatted Line objects.  Empty lines are not emitted.
1302
1303     Note: destroys the tree it's visiting by mutating prefixes of its leaves
1304     in ways that will no longer stringify to valid Python code on the tree.
1305     """
1306     is_pyi: bool = False
1307     current_line: Line = Factory(Line)
1308     remove_u_prefix: bool = False
1309
1310     def line(self, indent: int = 0, type: Type[Line] = Line) -> Iterator[Line]:
1311         """Generate a line.
1312
1313         If the line is empty, only emit if it makes sense.
1314         If the line is too long, split it first and then generate.
1315
1316         If any lines were generated, set up a new current_line.
1317         """
1318         if not self.current_line:
1319             if self.current_line.__class__ == type:
1320                 self.current_line.depth += indent
1321             else:
1322                 self.current_line = type(depth=self.current_line.depth + indent)
1323             return  # Line is empty, don't emit. Creating a new one unnecessary.
1324
1325         complete_line = self.current_line
1326         self.current_line = type(depth=complete_line.depth + indent)
1327         yield complete_line
1328
1329     def visit(self, node: LN) -> Iterator[Line]:
1330         """Main method to visit `node` and its children.
1331
1332         Yields :class:`Line` objects.
1333         """
1334         if isinstance(self.current_line, UnformattedLines):
1335             # File contained `# fmt: off`
1336             yield from self.visit_unformatted(node)
1337
1338         else:
1339             yield from super().visit(node)
1340
1341     def visit_default(self, node: LN) -> Iterator[Line]:
1342         """Default `visit_*()` implementation. Recurses to children of `node`."""
1343         if isinstance(node, Leaf):
1344             any_open_brackets = self.current_line.bracket_tracker.any_open_brackets()
1345             try:
1346                 for comment in generate_comments(node):
1347                     if any_open_brackets:
1348                         # any comment within brackets is subject to splitting
1349                         self.current_line.append(comment)
1350                     elif comment.type == token.COMMENT:
1351                         # regular trailing comment
1352                         self.current_line.append(comment)
1353                         yield from self.line()
1354
1355                     else:
1356                         # regular standalone comment
1357                         yield from self.line()
1358
1359                         self.current_line.append(comment)
1360                         yield from self.line()
1361
1362             except FormatOff as f_off:
1363                 f_off.trim_prefix(node)
1364                 yield from self.line(type=UnformattedLines)
1365                 yield from self.visit(node)
1366
1367             except FormatOn as f_on:
1368                 # This only happens here if somebody says "fmt: on" multiple
1369                 # times in a row.
1370                 f_on.trim_prefix(node)
1371                 yield from self.visit_default(node)
1372
1373             else:
1374                 normalize_prefix(node, inside_brackets=any_open_brackets)
1375                 if node.type == token.STRING:
1376                     normalize_string_prefix(node, remove_u_prefix=self.remove_u_prefix)
1377                     normalize_string_quotes(node)
1378                 if node.type not in WHITESPACE:
1379                     self.current_line.append(node)
1380         yield from super().visit_default(node)
1381
1382     def visit_INDENT(self, node: Node) -> Iterator[Line]:
1383         """Increase indentation level, maybe yield a line."""
1384         # In blib2to3 INDENT never holds comments.
1385         yield from self.line(+1)
1386         yield from self.visit_default(node)
1387
1388     def visit_DEDENT(self, node: Node) -> Iterator[Line]:
1389         """Decrease indentation level, maybe yield a line."""
1390         # The current line might still wait for trailing comments.  At DEDENT time
1391         # there won't be any (they would be prefixes on the preceding NEWLINE).
1392         # Emit the line then.
1393         yield from self.line()
1394
1395         # While DEDENT has no value, its prefix may contain standalone comments
1396         # that belong to the current indentation level.  Get 'em.
1397         yield from self.visit_default(node)
1398
1399         # Finally, emit the dedent.
1400         yield from self.line(-1)
1401
1402     def visit_stmt(
1403         self, node: Node, keywords: Set[str], parens: Set[str]
1404     ) -> Iterator[Line]:
1405         """Visit a statement.
1406
1407         This implementation is shared for `if`, `while`, `for`, `try`, `except`,
1408         `def`, `with`, `class`, `assert` and assignments.
1409
1410         The relevant Python language `keywords` for a given statement will be
1411         NAME leaves within it. This methods puts those on a separate line.
1412
1413         `parens` holds a set of string leaf values immediately after which
1414         invisible parens should be put.
1415         """
1416         normalize_invisible_parens(node, parens_after=parens)
1417         for child in node.children:
1418             if child.type == token.NAME and child.value in keywords:  # type: ignore
1419                 yield from self.line()
1420
1421             yield from self.visit(child)
1422
1423     def visit_suite(self, node: Node) -> Iterator[Line]:
1424         """Visit a suite."""
1425         if self.is_pyi and is_stub_suite(node):
1426             yield from self.visit(node.children[2])
1427         else:
1428             yield from self.visit_default(node)
1429
1430     def visit_simple_stmt(self, node: Node) -> Iterator[Line]:
1431         """Visit a statement without nested statements."""
1432         is_suite_like = node.parent and node.parent.type in STATEMENT
1433         if is_suite_like:
1434             if self.is_pyi and is_stub_body(node):
1435                 yield from self.visit_default(node)
1436             else:
1437                 yield from self.line(+1)
1438                 yield from self.visit_default(node)
1439                 yield from self.line(-1)
1440
1441         else:
1442             if not self.is_pyi or not node.parent or not is_stub_suite(node.parent):
1443                 yield from self.line()
1444             yield from self.visit_default(node)
1445
1446     def visit_async_stmt(self, node: Node) -> Iterator[Line]:
1447         """Visit `async def`, `async for`, `async with`."""
1448         yield from self.line()
1449
1450         children = iter(node.children)
1451         for child in children:
1452             yield from self.visit(child)
1453
1454             if child.type == token.ASYNC:
1455                 break
1456
1457         internal_stmt = next(children)
1458         for child in internal_stmt.children:
1459             yield from self.visit(child)
1460
1461     def visit_decorators(self, node: Node) -> Iterator[Line]:
1462         """Visit decorators."""
1463         for child in node.children:
1464             yield from self.line()
1465             yield from self.visit(child)
1466
1467     def visit_SEMI(self, leaf: Leaf) -> Iterator[Line]:
1468         """Remove a semicolon and put the other statement on a separate line."""
1469         yield from self.line()
1470
1471     def visit_ENDMARKER(self, leaf: Leaf) -> Iterator[Line]:
1472         """End of file. Process outstanding comments and end with a newline."""
1473         yield from self.visit_default(leaf)
1474         yield from self.line()
1475
1476     def visit_unformatted(self, node: LN) -> Iterator[Line]:
1477         """Used when file contained a `# fmt: off`."""
1478         if isinstance(node, Node):
1479             for child in node.children:
1480                 yield from self.visit(child)
1481
1482         else:
1483             try:
1484                 self.current_line.append(node)
1485             except FormatOn as f_on:
1486                 f_on.trim_prefix(node)
1487                 yield from self.line()
1488                 yield from self.visit(node)
1489
1490             if node.type == token.ENDMARKER:
1491                 # somebody decided not to put a final `# fmt: on`
1492                 yield from self.line()
1493
1494     def __attrs_post_init__(self) -> None:
1495         """You are in a twisty little maze of passages."""
1496         v = self.visit_stmt
1497         Ø: Set[str] = set()
1498         self.visit_assert_stmt = partial(v, keywords={"assert"}, parens={"assert", ","})
1499         self.visit_if_stmt = partial(
1500             v, keywords={"if", "else", "elif"}, parens={"if", "elif"}
1501         )
1502         self.visit_while_stmt = partial(v, keywords={"while", "else"}, parens={"while"})
1503         self.visit_for_stmt = partial(v, keywords={"for", "else"}, parens={"for", "in"})
1504         self.visit_try_stmt = partial(
1505             v, keywords={"try", "except", "else", "finally"}, parens=Ø
1506         )
1507         self.visit_except_clause = partial(v, keywords={"except"}, parens=Ø)
1508         self.visit_with_stmt = partial(v, keywords={"with"}, parens=Ø)
1509         self.visit_funcdef = partial(v, keywords={"def"}, parens=Ø)
1510         self.visit_classdef = partial(v, keywords={"class"}, parens=Ø)
1511         self.visit_expr_stmt = partial(v, keywords=Ø, parens=ASSIGNMENTS)
1512         self.visit_return_stmt = partial(v, keywords={"return"}, parens={"return"})
1513         self.visit_import_from = partial(v, keywords=Ø, parens={"import"})
1514         self.visit_async_funcdef = self.visit_async_stmt
1515         self.visit_decorated = self.visit_decorators
1516
1517
1518 IMPLICIT_TUPLE = {syms.testlist, syms.testlist_star_expr, syms.exprlist}
1519 BRACKET = {token.LPAR: token.RPAR, token.LSQB: token.RSQB, token.LBRACE: token.RBRACE}
1520 OPENING_BRACKETS = set(BRACKET.keys())
1521 CLOSING_BRACKETS = set(BRACKET.values())
1522 BRACKETS = OPENING_BRACKETS | CLOSING_BRACKETS
1523 ALWAYS_NO_SPACE = CLOSING_BRACKETS | {token.COMMA, STANDALONE_COMMENT}
1524
1525
1526 def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str:  # noqa C901
1527     """Return whitespace prefix if needed for the given `leaf`.
1528
1529     `complex_subscript` signals whether the given leaf is part of a subscription
1530     which has non-trivial arguments, like arithmetic expressions or function calls.
1531     """
1532     NO = ""
1533     SPACE = " "
1534     DOUBLESPACE = "  "
1535     t = leaf.type
1536     p = leaf.parent
1537     v = leaf.value
1538     if t in ALWAYS_NO_SPACE:
1539         return NO
1540
1541     if t == token.COMMENT:
1542         return DOUBLESPACE
1543
1544     assert p is not None, f"INTERNAL ERROR: hand-made leaf without parent: {leaf!r}"
1545     if t == token.COLON and p.type not in {
1546         syms.subscript,
1547         syms.subscriptlist,
1548         syms.sliceop,
1549     }:
1550         return NO
1551
1552     prev = leaf.prev_sibling
1553     if not prev:
1554         prevp = preceding_leaf(p)
1555         if not prevp or prevp.type in OPENING_BRACKETS:
1556             return NO
1557
1558         if t == token.COLON:
1559             if prevp.type == token.COLON:
1560                 return NO
1561
1562             elif prevp.type != token.COMMA and not complex_subscript:
1563                 return NO
1564
1565             return SPACE
1566
1567         if prevp.type == token.EQUAL:
1568             if prevp.parent:
1569                 if prevp.parent.type in {
1570                     syms.arglist,
1571                     syms.argument,
1572                     syms.parameters,
1573                     syms.varargslist,
1574                 }:
1575                     return NO
1576
1577                 elif prevp.parent.type == syms.typedargslist:
1578                     # A bit hacky: if the equal sign has whitespace, it means we
1579                     # previously found it's a typed argument.  So, we're using
1580                     # that, too.
1581                     return prevp.prefix
1582
1583         elif prevp.type in STARS:
1584             if is_vararg(prevp, within=VARARGS_PARENTS | UNPACKING_PARENTS):
1585                 return NO
1586
1587         elif prevp.type == token.COLON:
1588             if prevp.parent and prevp.parent.type in {syms.subscript, syms.sliceop}:
1589                 return SPACE if complex_subscript else NO
1590
1591         elif (
1592             prevp.parent
1593             and prevp.parent.type == syms.factor
1594             and prevp.type in MATH_OPERATORS
1595         ):
1596             return NO
1597
1598         elif (
1599             prevp.type == token.RIGHTSHIFT
1600             and prevp.parent
1601             and prevp.parent.type == syms.shift_expr
1602             and prevp.prev_sibling
1603             and prevp.prev_sibling.type == token.NAME
1604             and prevp.prev_sibling.value == "print"  # type: ignore
1605         ):
1606             # Python 2 print chevron
1607             return NO
1608
1609     elif prev.type in OPENING_BRACKETS:
1610         return NO
1611
1612     if p.type in {syms.parameters, syms.arglist}:
1613         # untyped function signatures or calls
1614         if not prev or prev.type != token.COMMA:
1615             return NO
1616
1617     elif p.type == syms.varargslist:
1618         # lambdas
1619         if prev and prev.type != token.COMMA:
1620             return NO
1621
1622     elif p.type == syms.typedargslist:
1623         # typed function signatures
1624         if not prev:
1625             return NO
1626
1627         if t == token.EQUAL:
1628             if prev.type != syms.tname:
1629                 return NO
1630
1631         elif prev.type == token.EQUAL:
1632             # A bit hacky: if the equal sign has whitespace, it means we
1633             # previously found it's a typed argument.  So, we're using that, too.
1634             return prev.prefix
1635
1636         elif prev.type != token.COMMA:
1637             return NO
1638
1639     elif p.type == syms.tname:
1640         # type names
1641         if not prev:
1642             prevp = preceding_leaf(p)
1643             if not prevp or prevp.type != token.COMMA:
1644                 return NO
1645
1646     elif p.type == syms.trailer:
1647         # attributes and calls
1648         if t == token.LPAR or t == token.RPAR:
1649             return NO
1650
1651         if not prev:
1652             if t == token.DOT:
1653                 prevp = preceding_leaf(p)
1654                 if not prevp or prevp.type != token.NUMBER:
1655                     return NO
1656
1657             elif t == token.LSQB:
1658                 return NO
1659
1660         elif prev.type != token.COMMA:
1661             return NO
1662
1663     elif p.type == syms.argument:
1664         # single argument
1665         if t == token.EQUAL:
1666             return NO
1667
1668         if not prev:
1669             prevp = preceding_leaf(p)
1670             if not prevp or prevp.type == token.LPAR:
1671                 return NO
1672
1673         elif prev.type in {token.EQUAL} | STARS:
1674             return NO
1675
1676     elif p.type == syms.decorator:
1677         # decorators
1678         return NO
1679
1680     elif p.type == syms.dotted_name:
1681         if prev:
1682             return NO
1683
1684         prevp = preceding_leaf(p)
1685         if not prevp or prevp.type == token.AT or prevp.type == token.DOT:
1686             return NO
1687
1688     elif p.type == syms.classdef:
1689         if t == token.LPAR:
1690             return NO
1691
1692         if prev and prev.type == token.LPAR:
1693             return NO
1694
1695     elif p.type in {syms.subscript, syms.sliceop}:
1696         # indexing
1697         if not prev:
1698             assert p.parent is not None, "subscripts are always parented"
1699             if p.parent.type == syms.subscriptlist:
1700                 return SPACE
1701
1702             return NO
1703
1704         elif not complex_subscript:
1705             return NO
1706
1707     elif p.type == syms.atom:
1708         if prev and t == token.DOT:
1709             # dots, but not the first one.
1710             return NO
1711
1712     elif p.type == syms.dictsetmaker:
1713         # dict unpacking
1714         if prev and prev.type == token.DOUBLESTAR:
1715             return NO
1716
1717     elif p.type in {syms.factor, syms.star_expr}:
1718         # unary ops
1719         if not prev:
1720             prevp = preceding_leaf(p)
1721             if not prevp or prevp.type in OPENING_BRACKETS:
1722                 return NO
1723
1724             prevp_parent = prevp.parent
1725             assert prevp_parent is not None
1726             if prevp.type == token.COLON and prevp_parent.type in {
1727                 syms.subscript,
1728                 syms.sliceop,
1729             }:
1730                 return NO
1731
1732             elif prevp.type == token.EQUAL and prevp_parent.type == syms.argument:
1733                 return NO
1734
1735         elif t == token.NAME or t == token.NUMBER:
1736             return NO
1737
1738     elif p.type == syms.import_from:
1739         if t == token.DOT:
1740             if prev and prev.type == token.DOT:
1741                 return NO
1742
1743         elif t == token.NAME:
1744             if v == "import":
1745                 return SPACE
1746
1747             if prev and prev.type == token.DOT:
1748                 return NO
1749
1750     elif p.type == syms.sliceop:
1751         return NO
1752
1753     return SPACE
1754
1755
1756 def preceding_leaf(node: Optional[LN]) -> Optional[Leaf]:
1757     """Return the first leaf that precedes `node`, if any."""
1758     while node:
1759         res = node.prev_sibling
1760         if res:
1761             if isinstance(res, Leaf):
1762                 return res
1763
1764             try:
1765                 return list(res.leaves())[-1]
1766
1767             except IndexError:
1768                 return None
1769
1770         node = node.parent
1771     return None
1772
1773
1774 def child_towards(ancestor: Node, descendant: LN) -> Optional[LN]:
1775     """Return the child of `ancestor` that contains `descendant`."""
1776     node: Optional[LN] = descendant
1777     while node and node.parent != ancestor:
1778         node = node.parent
1779     return node
1780
1781
1782 def is_split_after_delimiter(leaf: Leaf, previous: Leaf = None) -> int:
1783     """Return the priority of the `leaf` delimiter, given a line break after it.
1784
1785     The delimiter priorities returned here are from those delimiters that would
1786     cause a line break after themselves.
1787
1788     Higher numbers are higher priority.
1789     """
1790     if leaf.type == token.COMMA:
1791         return COMMA_PRIORITY
1792
1793     return 0
1794
1795
1796 def is_split_before_delimiter(leaf: Leaf, previous: Leaf = None) -> int:
1797     """Return the priority of the `leaf` delimiter, given a line before after it.
1798
1799     The delimiter priorities returned here are from those delimiters that would
1800     cause a line break before themselves.
1801
1802     Higher numbers are higher priority.
1803     """
1804     if is_vararg(leaf, within=VARARGS_PARENTS | UNPACKING_PARENTS):
1805         # * and ** might also be MATH_OPERATORS but in this case they are not.
1806         # Don't treat them as a delimiter.
1807         return 0
1808
1809     if (
1810         leaf.type == token.DOT
1811         and leaf.parent
1812         and leaf.parent.type not in {syms.import_from, syms.dotted_name}
1813         and (previous is None or previous.type in CLOSING_BRACKETS)
1814     ):
1815         return DOT_PRIORITY
1816
1817     if (
1818         leaf.type in MATH_OPERATORS
1819         and leaf.parent
1820         and leaf.parent.type not in {syms.factor, syms.star_expr}
1821     ):
1822         return MATH_PRIORITIES[leaf.type]
1823
1824     if leaf.type in COMPARATORS:
1825         return COMPARATOR_PRIORITY
1826
1827     if (
1828         leaf.type == token.STRING
1829         and previous is not None
1830         and previous.type == token.STRING
1831     ):
1832         return STRING_PRIORITY
1833
1834     if leaf.type != token.NAME:
1835         return 0
1836
1837     if (
1838         leaf.value == "for"
1839         and leaf.parent
1840         and leaf.parent.type in {syms.comp_for, syms.old_comp_for}
1841     ):
1842         return COMPREHENSION_PRIORITY
1843
1844     if (
1845         leaf.value == "if"
1846         and leaf.parent
1847         and leaf.parent.type in {syms.comp_if, syms.old_comp_if}
1848     ):
1849         return COMPREHENSION_PRIORITY
1850
1851     if leaf.value in {"if", "else"} and leaf.parent and leaf.parent.type == syms.test:
1852         return TERNARY_PRIORITY
1853
1854     if leaf.value == "is":
1855         return COMPARATOR_PRIORITY
1856
1857     if (
1858         leaf.value == "in"
1859         and leaf.parent
1860         and leaf.parent.type in {syms.comp_op, syms.comparison}
1861         and not (
1862             previous is not None
1863             and previous.type == token.NAME
1864             and previous.value == "not"
1865         )
1866     ):
1867         return COMPARATOR_PRIORITY
1868
1869     if (
1870         leaf.value == "not"
1871         and leaf.parent
1872         and leaf.parent.type == syms.comp_op
1873         and not (
1874             previous is not None
1875             and previous.type == token.NAME
1876             and previous.value == "is"
1877         )
1878     ):
1879         return COMPARATOR_PRIORITY
1880
1881     if leaf.value in LOGIC_OPERATORS and leaf.parent:
1882         return LOGIC_PRIORITY
1883
1884     return 0
1885
1886
1887 def generate_comments(leaf: LN) -> Iterator[Leaf]:
1888     """Clean the prefix of the `leaf` and generate comments from it, if any.
1889
1890     Comments in lib2to3 are shoved into the whitespace prefix.  This happens
1891     in `pgen2/driver.py:Driver.parse_tokens()`.  This was a brilliant implementation
1892     move because it does away with modifying the grammar to include all the
1893     possible places in which comments can be placed.
1894
1895     The sad consequence for us though is that comments don't "belong" anywhere.
1896     This is why this function generates simple parentless Leaf objects for
1897     comments.  We simply don't know what the correct parent should be.
1898
1899     No matter though, we can live without this.  We really only need to
1900     differentiate between inline and standalone comments.  The latter don't
1901     share the line with any code.
1902
1903     Inline comments are emitted as regular token.COMMENT leaves.  Standalone
1904     are emitted with a fake STANDALONE_COMMENT token identifier.
1905     """
1906     p = leaf.prefix
1907     if not p:
1908         return
1909
1910     if "#" not in p:
1911         return
1912
1913     consumed = 0
1914     nlines = 0
1915     for index, line in enumerate(p.split("\n")):
1916         consumed += len(line) + 1  # adding the length of the split '\n'
1917         line = line.lstrip()
1918         if not line:
1919             nlines += 1
1920         if not line.startswith("#"):
1921             continue
1922
1923         if index == 0 and leaf.type != token.ENDMARKER:
1924             comment_type = token.COMMENT  # simple trailing comment
1925         else:
1926             comment_type = STANDALONE_COMMENT
1927         comment = make_comment(line)
1928         yield Leaf(comment_type, comment, prefix="\n" * nlines)
1929
1930         if comment in {"# fmt: on", "# yapf: enable"}:
1931             raise FormatOn(consumed)
1932
1933         if comment in {"# fmt: off", "# yapf: disable"}:
1934             if comment_type == STANDALONE_COMMENT:
1935                 raise FormatOff(consumed)
1936
1937             prev = preceding_leaf(leaf)
1938             if not prev or prev.type in WHITESPACE:  # standalone comment in disguise
1939                 raise FormatOff(consumed)
1940
1941         nlines = 0
1942
1943
1944 def make_comment(content: str) -> str:
1945     """Return a consistently formatted comment from the given `content` string.
1946
1947     All comments (except for "##", "#!", "#:") should have a single space between
1948     the hash sign and the content.
1949
1950     If `content` didn't start with a hash sign, one is provided.
1951     """
1952     content = content.rstrip()
1953     if not content:
1954         return "#"
1955
1956     if content[0] == "#":
1957         content = content[1:]
1958     if content and content[0] not in " !:#":
1959         content = " " + content
1960     return "#" + content
1961
1962
1963 def split_line(
1964     line: Line, line_length: int, inner: bool = False, py36: bool = False
1965 ) -> Iterator[Line]:
1966     """Split a `line` into potentially many lines.
1967
1968     They should fit in the allotted `line_length` but might not be able to.
1969     `inner` signifies that there were a pair of brackets somewhere around the
1970     current `line`, possibly transitively. This means we can fallback to splitting
1971     by delimiters if the LHS/RHS don't yield any results.
1972
1973     If `py36` is True, splitting may generate syntax that is only compatible
1974     with Python 3.6 and later.
1975     """
1976     if isinstance(line, UnformattedLines) or line.is_comment:
1977         yield line
1978         return
1979
1980     line_str = str(line).strip("\n")
1981     if not line.should_explode and is_line_short_enough(
1982         line, line_length=line_length, line_str=line_str
1983     ):
1984         yield line
1985         return
1986
1987     split_funcs: List[SplitFunc]
1988     if line.is_def:
1989         split_funcs = [left_hand_split]
1990     else:
1991
1992         def rhs(line: Line, py36: bool = False) -> Iterator[Line]:
1993             for omit in generate_trailers_to_omit(line, line_length):
1994                 lines = list(right_hand_split(line, line_length, py36, omit=omit))
1995                 if is_line_short_enough(lines[0], line_length=line_length):
1996                     yield from lines
1997                     return
1998
1999             # All splits failed, best effort split with no omits.
2000             # This mostly happens to multiline strings that are by definition
2001             # reported as not fitting a single line.
2002             yield from right_hand_split(line, py36)
2003
2004         if line.inside_brackets:
2005             split_funcs = [delimiter_split, standalone_comment_split, rhs]
2006         else:
2007             split_funcs = [rhs]
2008     for split_func in split_funcs:
2009         # We are accumulating lines in `result` because we might want to abort
2010         # mission and return the original line in the end, or attempt a different
2011         # split altogether.
2012         result: List[Line] = []
2013         try:
2014             for l in split_func(line, py36):
2015                 if str(l).strip("\n") == line_str:
2016                     raise CannotSplit("Split function returned an unchanged result")
2017
2018                 result.extend(
2019                     split_line(l, line_length=line_length, inner=True, py36=py36)
2020                 )
2021         except CannotSplit as cs:
2022             continue
2023
2024         else:
2025             yield from result
2026             break
2027
2028     else:
2029         yield line
2030
2031
2032 def left_hand_split(line: Line, py36: bool = False) -> Iterator[Line]:
2033     """Split line into many lines, starting with the first matching bracket pair.
2034
2035     Note: this usually looks weird, only use this for function definitions.
2036     Prefer RHS otherwise.  This is why this function is not symmetrical with
2037     :func:`right_hand_split` which also handles optional parentheses.
2038     """
2039     head = Line(depth=line.depth)
2040     body = Line(depth=line.depth + 1, inside_brackets=True)
2041     tail = Line(depth=line.depth)
2042     tail_leaves: List[Leaf] = []
2043     body_leaves: List[Leaf] = []
2044     head_leaves: List[Leaf] = []
2045     current_leaves = head_leaves
2046     matching_bracket = None
2047     for leaf in line.leaves:
2048         if (
2049             current_leaves is body_leaves
2050             and leaf.type in CLOSING_BRACKETS
2051             and leaf.opening_bracket is matching_bracket
2052         ):
2053             current_leaves = tail_leaves if body_leaves else head_leaves
2054         current_leaves.append(leaf)
2055         if current_leaves is head_leaves:
2056             if leaf.type in OPENING_BRACKETS:
2057                 matching_bracket = leaf
2058                 current_leaves = body_leaves
2059     # Since body is a new indent level, remove spurious leading whitespace.
2060     if body_leaves:
2061         normalize_prefix(body_leaves[0], inside_brackets=True)
2062     # Build the new lines.
2063     for result, leaves in (head, head_leaves), (body, body_leaves), (tail, tail_leaves):
2064         for leaf in leaves:
2065             result.append(leaf, preformatted=True)
2066             for comment_after in line.comments_after(leaf):
2067                 result.append(comment_after, preformatted=True)
2068     bracket_split_succeeded_or_raise(head, body, tail)
2069     for result in (head, body, tail):
2070         if result:
2071             yield result
2072
2073
2074 def right_hand_split(
2075     line: Line, line_length: int, py36: bool = False, omit: Collection[LeafID] = ()
2076 ) -> Iterator[Line]:
2077     """Split line into many lines, starting with the last matching bracket pair.
2078
2079     If the split was by optional parentheses, attempt splitting without them, too.
2080     `omit` is a collection of closing bracket IDs that shouldn't be considered for
2081     this split.
2082
2083     Note: running this function modifies `bracket_depth` on the leaves of `line`.
2084     """
2085     head = Line(depth=line.depth)
2086     body = Line(depth=line.depth + 1, inside_brackets=True)
2087     tail = Line(depth=line.depth)
2088     tail_leaves: List[Leaf] = []
2089     body_leaves: List[Leaf] = []
2090     head_leaves: List[Leaf] = []
2091     current_leaves = tail_leaves
2092     opening_bracket = None
2093     closing_bracket = None
2094     for leaf in reversed(line.leaves):
2095         if current_leaves is body_leaves:
2096             if leaf is opening_bracket:
2097                 current_leaves = head_leaves if body_leaves else tail_leaves
2098         current_leaves.append(leaf)
2099         if current_leaves is tail_leaves:
2100             if leaf.type in CLOSING_BRACKETS and id(leaf) not in omit:
2101                 opening_bracket = leaf.opening_bracket
2102                 closing_bracket = leaf
2103                 current_leaves = body_leaves
2104     tail_leaves.reverse()
2105     body_leaves.reverse()
2106     head_leaves.reverse()
2107     # Since body is a new indent level, remove spurious leading whitespace.
2108     if body_leaves:
2109         normalize_prefix(body_leaves[0], inside_brackets=True)
2110     if not head_leaves:
2111         # No `head` means the split failed. Either `tail` has all content or
2112         # the matching `opening_bracket` wasn't available on `line` anymore.
2113         raise CannotSplit("No brackets found")
2114
2115     # Build the new lines.
2116     for result, leaves in (head, head_leaves), (body, body_leaves), (tail, tail_leaves):
2117         for leaf in leaves:
2118             result.append(leaf, preformatted=True)
2119             for comment_after in line.comments_after(leaf):
2120                 result.append(comment_after, preformatted=True)
2121     bracket_split_succeeded_or_raise(head, body, tail)
2122     assert opening_bracket and closing_bracket
2123     if (
2124         # the opening bracket is an optional paren
2125         opening_bracket.type == token.LPAR
2126         and not opening_bracket.value
2127         # the closing bracket is an optional paren
2128         and closing_bracket.type == token.RPAR
2129         and not closing_bracket.value
2130         # there are no standalone comments in the body
2131         and not line.contains_standalone_comments(0)
2132         # and it's not an import (optional parens are the only thing we can split
2133         # on in this case; attempting a split without them is a waste of time)
2134         and not line.is_import
2135     ):
2136         omit = {id(closing_bracket), *omit}
2137         if can_omit_invisible_parens(body, line_length):
2138             try:
2139                 yield from right_hand_split(line, line_length, py36=py36, omit=omit)
2140                 return
2141             except CannotSplit:
2142                 pass
2143
2144     ensure_visible(opening_bracket)
2145     ensure_visible(closing_bracket)
2146     body.should_explode = should_explode(body, opening_bracket)
2147     for result in (head, body, tail):
2148         if result:
2149             yield result
2150
2151
2152 def bracket_split_succeeded_or_raise(head: Line, body: Line, tail: Line) -> None:
2153     """Raise :exc:`CannotSplit` if the last left- or right-hand split failed.
2154
2155     Do nothing otherwise.
2156
2157     A left- or right-hand split is based on a pair of brackets. Content before
2158     (and including) the opening bracket is left on one line, content inside the
2159     brackets is put on a separate line, and finally content starting with and
2160     following the closing bracket is put on a separate line.
2161
2162     Those are called `head`, `body`, and `tail`, respectively. If the split
2163     produced the same line (all content in `head`) or ended up with an empty `body`
2164     and the `tail` is just the closing bracket, then it's considered failed.
2165     """
2166     tail_len = len(str(tail).strip())
2167     if not body:
2168         if tail_len == 0:
2169             raise CannotSplit("Splitting brackets produced the same line")
2170
2171         elif tail_len < 3:
2172             raise CannotSplit(
2173                 f"Splitting brackets on an empty body to save "
2174                 f"{tail_len} characters is not worth it"
2175             )
2176
2177
2178 def dont_increase_indentation(split_func: SplitFunc) -> SplitFunc:
2179     """Normalize prefix of the first leaf in every line returned by `split_func`.
2180
2181     This is a decorator over relevant split functions.
2182     """
2183
2184     @wraps(split_func)
2185     def split_wrapper(line: Line, py36: bool = False) -> Iterator[Line]:
2186         for l in split_func(line, py36):
2187             normalize_prefix(l.leaves[0], inside_brackets=True)
2188             yield l
2189
2190     return split_wrapper
2191
2192
2193 @dont_increase_indentation
2194 def delimiter_split(line: Line, py36: bool = False) -> Iterator[Line]:
2195     """Split according to delimiters of the highest priority.
2196
2197     If `py36` is True, the split will add trailing commas also in function
2198     signatures that contain `*` and `**`.
2199     """
2200     try:
2201         last_leaf = line.leaves[-1]
2202     except IndexError:
2203         raise CannotSplit("Line empty")
2204
2205     bt = line.bracket_tracker
2206     try:
2207         delimiter_priority = bt.max_delimiter_priority(exclude={id(last_leaf)})
2208     except ValueError:
2209         raise CannotSplit("No delimiters found")
2210
2211     if delimiter_priority == DOT_PRIORITY:
2212         if bt.delimiter_count_with_priority(delimiter_priority) == 1:
2213             raise CannotSplit("Splitting a single attribute from its owner looks wrong")
2214
2215     current_line = Line(depth=line.depth, inside_brackets=line.inside_brackets)
2216     lowest_depth = sys.maxsize
2217     trailing_comma_safe = True
2218
2219     def append_to_line(leaf: Leaf) -> Iterator[Line]:
2220         """Append `leaf` to current line or to new line if appending impossible."""
2221         nonlocal current_line
2222         try:
2223             current_line.append_safe(leaf, preformatted=True)
2224         except ValueError as ve:
2225             yield current_line
2226
2227             current_line = Line(depth=line.depth, inside_brackets=line.inside_brackets)
2228             current_line.append(leaf)
2229
2230     for index, leaf in enumerate(line.leaves):
2231         yield from append_to_line(leaf)
2232
2233         for comment_after in line.comments_after(leaf, index):
2234             yield from append_to_line(comment_after)
2235
2236         lowest_depth = min(lowest_depth, leaf.bracket_depth)
2237         if leaf.bracket_depth == lowest_depth and is_vararg(
2238             leaf, within=VARARGS_PARENTS
2239         ):
2240             trailing_comma_safe = trailing_comma_safe and py36
2241         leaf_priority = bt.delimiters.get(id(leaf))
2242         if leaf_priority == delimiter_priority:
2243             yield current_line
2244
2245             current_line = Line(depth=line.depth, inside_brackets=line.inside_brackets)
2246     if current_line:
2247         if (
2248             trailing_comma_safe
2249             and delimiter_priority == COMMA_PRIORITY
2250             and current_line.leaves[-1].type != token.COMMA
2251             and current_line.leaves[-1].type != STANDALONE_COMMENT
2252         ):
2253             current_line.append(Leaf(token.COMMA, ","))
2254         yield current_line
2255
2256
2257 @dont_increase_indentation
2258 def standalone_comment_split(line: Line, py36: bool = False) -> Iterator[Line]:
2259     """Split standalone comments from the rest of the line."""
2260     if not line.contains_standalone_comments(0):
2261         raise CannotSplit("Line does not have any standalone comments")
2262
2263     current_line = Line(depth=line.depth, inside_brackets=line.inside_brackets)
2264
2265     def append_to_line(leaf: Leaf) -> Iterator[Line]:
2266         """Append `leaf` to current line or to new line if appending impossible."""
2267         nonlocal current_line
2268         try:
2269             current_line.append_safe(leaf, preformatted=True)
2270         except ValueError as ve:
2271             yield current_line
2272
2273             current_line = Line(depth=line.depth, inside_brackets=line.inside_brackets)
2274             current_line.append(leaf)
2275
2276     for index, leaf in enumerate(line.leaves):
2277         yield from append_to_line(leaf)
2278
2279         for comment_after in line.comments_after(leaf, index):
2280             yield from append_to_line(comment_after)
2281
2282     if current_line:
2283         yield current_line
2284
2285
2286 def is_import(leaf: Leaf) -> bool:
2287     """Return True if the given leaf starts an import statement."""
2288     p = leaf.parent
2289     t = leaf.type
2290     v = leaf.value
2291     return bool(
2292         t == token.NAME
2293         and (
2294             (v == "import" and p and p.type == syms.import_name)
2295             or (v == "from" and p and p.type == syms.import_from)
2296         )
2297     )
2298
2299
2300 def normalize_prefix(leaf: Leaf, *, inside_brackets: bool) -> None:
2301     """Leave existing extra newlines if not `inside_brackets`. Remove everything
2302     else.
2303
2304     Note: don't use backslashes for formatting or you'll lose your voting rights.
2305     """
2306     if not inside_brackets:
2307         spl = leaf.prefix.split("#")
2308         if "\\" not in spl[0]:
2309             nl_count = spl[-1].count("\n")
2310             if len(spl) > 1:
2311                 nl_count -= 1
2312             leaf.prefix = "\n" * nl_count
2313             return
2314
2315     leaf.prefix = ""
2316
2317
2318 def normalize_string_prefix(leaf: Leaf, remove_u_prefix: bool = False) -> None:
2319     """Make all string prefixes lowercase.
2320
2321     If remove_u_prefix is given, also removes any u prefix from the string.
2322
2323     Note: Mutates its argument.
2324     """
2325     match = re.match(r"^([furbFURB]*)(.*)$", leaf.value, re.DOTALL)
2326     assert match is not None, f"failed to match string {leaf.value!r}"
2327     orig_prefix = match.group(1)
2328     new_prefix = orig_prefix.lower()
2329     if remove_u_prefix:
2330         new_prefix = new_prefix.replace("u", "")
2331     leaf.value = f"{new_prefix}{match.group(2)}"
2332
2333
2334 def normalize_string_quotes(leaf: Leaf) -> None:
2335     """Prefer double quotes but only if it doesn't cause more escaping.
2336
2337     Adds or removes backslashes as appropriate. Doesn't parse and fix
2338     strings nested in f-strings (yet).
2339
2340     Note: Mutates its argument.
2341     """
2342     value = leaf.value.lstrip("furbFURB")
2343     if value[:3] == '"""':
2344         return
2345
2346     elif value[:3] == "'''":
2347         orig_quote = "'''"
2348         new_quote = '"""'
2349     elif value[0] == '"':
2350         orig_quote = '"'
2351         new_quote = "'"
2352     else:
2353         orig_quote = "'"
2354         new_quote = '"'
2355     first_quote_pos = leaf.value.find(orig_quote)
2356     if first_quote_pos == -1:
2357         return  # There's an internal error
2358
2359     prefix = leaf.value[:first_quote_pos]
2360     unescaped_new_quote = re.compile(rf"(([^\\]|^)(\\\\)*){new_quote}")
2361     escaped_new_quote = re.compile(rf"([^\\]|^)\\(\\\\)*{new_quote}")
2362     escaped_orig_quote = re.compile(rf"([^\\]|^)\\(\\\\)*{orig_quote}")
2363     body = leaf.value[first_quote_pos + len(orig_quote) : -len(orig_quote)]
2364     if "r" in prefix.casefold():
2365         if unescaped_new_quote.search(body):
2366             # There's at least one unescaped new_quote in this raw string
2367             # so converting is impossible
2368             return
2369
2370         # Do not introduce or remove backslashes in raw strings
2371         new_body = body
2372     else:
2373         # remove unnecessary quotes
2374         new_body = sub_twice(escaped_new_quote, rf"\1\2{new_quote}", body)
2375         if body != new_body:
2376             # Consider the string without unnecessary quotes as the original
2377             body = new_body
2378             leaf.value = f"{prefix}{orig_quote}{body}{orig_quote}"
2379         new_body = sub_twice(escaped_orig_quote, rf"\1\2{orig_quote}", new_body)
2380         new_body = sub_twice(unescaped_new_quote, rf"\1\\{new_quote}", new_body)
2381     if new_quote == '"""' and new_body[-1] == '"':
2382         # edge case:
2383         new_body = new_body[:-1] + '\\"'
2384     orig_escape_count = body.count("\\")
2385     new_escape_count = new_body.count("\\")
2386     if new_escape_count > orig_escape_count:
2387         return  # Do not introduce more escaping
2388
2389     if new_escape_count == orig_escape_count and orig_quote == '"':
2390         return  # Prefer double quotes
2391
2392     leaf.value = f"{prefix}{new_quote}{new_body}{new_quote}"
2393
2394
2395 def normalize_invisible_parens(node: Node, parens_after: Set[str]) -> None:
2396     """Make existing optional parentheses invisible or create new ones.
2397
2398     `parens_after` is a set of string leaf values immeditely after which parens
2399     should be put.
2400
2401     Standardizes on visible parentheses for single-element tuples, and keeps
2402     existing visible parentheses for other tuples and generator expressions.
2403     """
2404     try:
2405         list(generate_comments(node))
2406     except FormatOff:
2407         return  # This `node` has a prefix with `# fmt: off`, don't mess with parens.
2408
2409     check_lpar = False
2410     for index, child in enumerate(list(node.children)):
2411         if check_lpar:
2412             if child.type == syms.atom:
2413                 maybe_make_parens_invisible_in_atom(child)
2414             elif is_one_tuple(child):
2415                 # wrap child in visible parentheses
2416                 lpar = Leaf(token.LPAR, "(")
2417                 rpar = Leaf(token.RPAR, ")")
2418                 child.remove()
2419                 node.insert_child(index, Node(syms.atom, [lpar, child, rpar]))
2420             elif node.type == syms.import_from:
2421                 # "import from" nodes store parentheses directly as part of
2422                 # the statement
2423                 if child.type == token.LPAR:
2424                     # make parentheses invisible
2425                     child.value = ""  # type: ignore
2426                     node.children[-1].value = ""  # type: ignore
2427                 elif child.type != token.STAR:
2428                     # insert invisible parentheses
2429                     node.insert_child(index, Leaf(token.LPAR, ""))
2430                     node.append_child(Leaf(token.RPAR, ""))
2431                 break
2432
2433             elif not (isinstance(child, Leaf) and is_multiline_string(child)):
2434                 # wrap child in invisible parentheses
2435                 lpar = Leaf(token.LPAR, "")
2436                 rpar = Leaf(token.RPAR, "")
2437                 index = child.remove() or 0
2438                 node.insert_child(index, Node(syms.atom, [lpar, child, rpar]))
2439
2440         check_lpar = isinstance(child, Leaf) and child.value in parens_after
2441
2442
2443 def maybe_make_parens_invisible_in_atom(node: LN) -> bool:
2444     """If it's safe, make the parens in the atom `node` invisible, recusively."""
2445     if (
2446         node.type != syms.atom
2447         or is_empty_tuple(node)
2448         or is_one_tuple(node)
2449         or is_yield(node)
2450         or max_delimiter_priority_in_atom(node) >= COMMA_PRIORITY
2451     ):
2452         return False
2453
2454     first = node.children[0]
2455     last = node.children[-1]
2456     if first.type == token.LPAR and last.type == token.RPAR:
2457         # make parentheses invisible
2458         first.value = ""  # type: ignore
2459         last.value = ""  # type: ignore
2460         if len(node.children) > 1:
2461             maybe_make_parens_invisible_in_atom(node.children[1])
2462         return True
2463
2464     return False
2465
2466
2467 def is_empty_tuple(node: LN) -> bool:
2468     """Return True if `node` holds an empty tuple."""
2469     return (
2470         node.type == syms.atom
2471         and len(node.children) == 2
2472         and node.children[0].type == token.LPAR
2473         and node.children[1].type == token.RPAR
2474     )
2475
2476
2477 def is_one_tuple(node: LN) -> bool:
2478     """Return True if `node` holds a tuple with one element, with or without parens."""
2479     if node.type == syms.atom:
2480         if len(node.children) != 3:
2481             return False
2482
2483         lpar, gexp, rpar = node.children
2484         if not (
2485             lpar.type == token.LPAR
2486             and gexp.type == syms.testlist_gexp
2487             and rpar.type == token.RPAR
2488         ):
2489             return False
2490
2491         return len(gexp.children) == 2 and gexp.children[1].type == token.COMMA
2492
2493     return (
2494         node.type in IMPLICIT_TUPLE
2495         and len(node.children) == 2
2496         and node.children[1].type == token.COMMA
2497     )
2498
2499
2500 def is_yield(node: LN) -> bool:
2501     """Return True if `node` holds a `yield` or `yield from` expression."""
2502     if node.type == syms.yield_expr:
2503         return True
2504
2505     if node.type == token.NAME and node.value == "yield":  # type: ignore
2506         return True
2507
2508     if node.type != syms.atom:
2509         return False
2510
2511     if len(node.children) != 3:
2512         return False
2513
2514     lpar, expr, rpar = node.children
2515     if lpar.type == token.LPAR and rpar.type == token.RPAR:
2516         return is_yield(expr)
2517
2518     return False
2519
2520
2521 def is_vararg(leaf: Leaf, within: Set[NodeType]) -> bool:
2522     """Return True if `leaf` is a star or double star in a vararg or kwarg.
2523
2524     If `within` includes VARARGS_PARENTS, this applies to function signatures.
2525     If `within` includes UNPACKING_PARENTS, it applies to right hand-side
2526     extended iterable unpacking (PEP 3132) and additional unpacking
2527     generalizations (PEP 448).
2528     """
2529     if leaf.type not in STARS or not leaf.parent:
2530         return False
2531
2532     p = leaf.parent
2533     if p.type == syms.star_expr:
2534         # Star expressions are also used as assignment targets in extended
2535         # iterable unpacking (PEP 3132).  See what its parent is instead.
2536         if not p.parent:
2537             return False
2538
2539         p = p.parent
2540
2541     return p.type in within
2542
2543
2544 def is_multiline_string(leaf: Leaf) -> bool:
2545     """Return True if `leaf` is a multiline string that actually spans many lines."""
2546     value = leaf.value.lstrip("furbFURB")
2547     return value[:3] in {'"""', "'''"} and "\n" in value
2548
2549
2550 def is_stub_suite(node: Node) -> bool:
2551     """Return True if `node` is a suite with a stub body."""
2552     if (
2553         len(node.children) != 4
2554         or node.children[0].type != token.NEWLINE
2555         or node.children[1].type != token.INDENT
2556         or node.children[3].type != token.DEDENT
2557     ):
2558         return False
2559
2560     return is_stub_body(node.children[2])
2561
2562
2563 def is_stub_body(node: LN) -> bool:
2564     """Return True if `node` is a simple statement containing an ellipsis."""
2565     if not isinstance(node, Node) or node.type != syms.simple_stmt:
2566         return False
2567
2568     if len(node.children) != 2:
2569         return False
2570
2571     child = node.children[0]
2572     return (
2573         child.type == syms.atom
2574         and len(child.children) == 3
2575         and all(leaf == Leaf(token.DOT, ".") for leaf in child.children)
2576     )
2577
2578
2579 def max_delimiter_priority_in_atom(node: LN) -> int:
2580     """Return maximum delimiter priority inside `node`.
2581
2582     This is specific to atoms with contents contained in a pair of parentheses.
2583     If `node` isn't an atom or there are no enclosing parentheses, returns 0.
2584     """
2585     if node.type != syms.atom:
2586         return 0
2587
2588     first = node.children[0]
2589     last = node.children[-1]
2590     if not (first.type == token.LPAR and last.type == token.RPAR):
2591         return 0
2592
2593     bt = BracketTracker()
2594     for c in node.children[1:-1]:
2595         if isinstance(c, Leaf):
2596             bt.mark(c)
2597         else:
2598             for leaf in c.leaves():
2599                 bt.mark(leaf)
2600     try:
2601         return bt.max_delimiter_priority()
2602
2603     except ValueError:
2604         return 0
2605
2606
2607 def ensure_visible(leaf: Leaf) -> None:
2608     """Make sure parentheses are visible.
2609
2610     They could be invisible as part of some statements (see
2611     :func:`normalize_invible_parens` and :func:`visit_import_from`).
2612     """
2613     if leaf.type == token.LPAR:
2614         leaf.value = "("
2615     elif leaf.type == token.RPAR:
2616         leaf.value = ")"
2617
2618
2619 def should_explode(line: Line, opening_bracket: Leaf) -> bool:
2620     """Should `line` immediately be split with `delimiter_split()` after RHS?"""
2621     if not (
2622         opening_bracket.parent
2623         and opening_bracket.parent.type in {syms.atom, syms.import_from}
2624         and opening_bracket.value in "[{("
2625     ):
2626         return False
2627
2628     try:
2629         last_leaf = line.leaves[-1]
2630         exclude = {id(last_leaf)} if last_leaf.type == token.COMMA else set()
2631         max_priority = line.bracket_tracker.max_delimiter_priority(exclude=exclude)
2632     except (IndexError, ValueError):
2633         return False
2634
2635     return max_priority == COMMA_PRIORITY
2636
2637
2638 def is_python36(node: Node) -> bool:
2639     """Return True if the current file is using Python 3.6+ features.
2640
2641     Currently looking for:
2642     - f-strings; and
2643     - trailing commas after * or ** in function signatures and calls.
2644     """
2645     for n in node.pre_order():
2646         if n.type == token.STRING:
2647             value_head = n.value[:2]  # type: ignore
2648             if value_head in {'f"', 'F"', "f'", "F'", "rf", "fr", "RF", "FR"}:
2649                 return True
2650
2651         elif (
2652             n.type in {syms.typedargslist, syms.arglist}
2653             and n.children
2654             and n.children[-1].type == token.COMMA
2655         ):
2656             for ch in n.children:
2657                 if ch.type in STARS:
2658                     return True
2659
2660                 if ch.type == syms.argument:
2661                     for argch in ch.children:
2662                         if argch.type in STARS:
2663                             return True
2664
2665     return False
2666
2667
2668 def generate_trailers_to_omit(line: Line, line_length: int) -> Iterator[Set[LeafID]]:
2669     """Generate sets of closing bracket IDs that should be omitted in a RHS.
2670
2671     Brackets can be omitted if the entire trailer up to and including
2672     a preceding closing bracket fits in one line.
2673
2674     Yielded sets are cumulative (contain results of previous yields, too).  First
2675     set is empty.
2676     """
2677
2678     omit: Set[LeafID] = set()
2679     yield omit
2680
2681     length = 4 * line.depth
2682     opening_bracket = None
2683     closing_bracket = None
2684     optional_brackets: Set[LeafID] = set()
2685     inner_brackets: Set[LeafID] = set()
2686     for index, leaf, leaf_length in enumerate_with_length(line, reversed=True):
2687         length += leaf_length
2688         if length > line_length:
2689             break
2690
2691         has_inline_comment = leaf_length > len(leaf.value) + len(leaf.prefix)
2692         if leaf.type == STANDALONE_COMMENT or has_inline_comment:
2693             break
2694
2695         optional_brackets.discard(id(leaf))
2696         if opening_bracket:
2697             if leaf is opening_bracket:
2698                 opening_bracket = None
2699             elif leaf.type in CLOSING_BRACKETS:
2700                 inner_brackets.add(id(leaf))
2701         elif leaf.type in CLOSING_BRACKETS:
2702             if not leaf.value:
2703                 optional_brackets.add(id(opening_bracket))
2704                 continue
2705
2706             if index > 0 and line.leaves[index - 1].type in OPENING_BRACKETS:
2707                 # Empty brackets would fail a split so treat them as "inner"
2708                 # brackets (e.g. only add them to the `omit` set if another
2709                 # pair of brackets was good enough.
2710                 inner_brackets.add(id(leaf))
2711                 continue
2712
2713             opening_bracket = leaf.opening_bracket
2714             if closing_bracket:
2715                 omit.add(id(closing_bracket))
2716                 omit.update(inner_brackets)
2717                 inner_brackets.clear()
2718                 yield omit
2719             closing_bracket = leaf
2720
2721
2722 def get_future_imports(node: Node) -> Set[str]:
2723     """Return a set of __future__ imports in the file."""
2724     imports = set()
2725     for child in node.children:
2726         if child.type != syms.simple_stmt:
2727             break
2728         first_child = child.children[0]
2729         if isinstance(first_child, Leaf):
2730             # Continue looking if we see a docstring; otherwise stop.
2731             if (
2732                 len(child.children) == 2
2733                 and first_child.type == token.STRING
2734                 and child.children[1].type == token.NEWLINE
2735             ):
2736                 continue
2737             else:
2738                 break
2739         elif first_child.type == syms.import_from:
2740             module_name = first_child.children[1]
2741             if not isinstance(module_name, Leaf) or module_name.value != "__future__":
2742                 break
2743             for import_from_child in first_child.children[3:]:
2744                 if isinstance(import_from_child, Leaf):
2745                     if import_from_child.type == token.NAME:
2746                         imports.add(import_from_child.value)
2747                 else:
2748                     assert import_from_child.type == syms.import_as_names
2749                     for leaf in import_from_child.children:
2750                         if isinstance(leaf, Leaf) and leaf.type == token.NAME:
2751                             imports.add(leaf.value)
2752         else:
2753             break
2754     return imports
2755
2756
2757 PYTHON_EXTENSIONS = {".py", ".pyi"}
2758 BLACKLISTED_DIRECTORIES = {
2759     "build",
2760     "buck-out",
2761     "dist",
2762     "_build",
2763     ".git",
2764     ".hg",
2765     ".mypy_cache",
2766     ".tox",
2767     ".venv",
2768 }
2769
2770
2771 def gen_python_files_in_dir(path: Path) -> Iterator[Path]:
2772     """Generate all files under `path` which aren't under BLACKLISTED_DIRECTORIES
2773     and have one of the PYTHON_EXTENSIONS.
2774     """
2775     for child in path.iterdir():
2776         if child.is_dir():
2777             if child.name in BLACKLISTED_DIRECTORIES:
2778                 continue
2779
2780             yield from gen_python_files_in_dir(child)
2781
2782         elif child.is_file() and child.suffix in PYTHON_EXTENSIONS:
2783             yield child
2784
2785
2786 @dataclass
2787 class Report:
2788     """Provides a reformatting counter. Can be rendered with `str(report)`."""
2789     check: bool = False
2790     quiet: bool = False
2791     change_count: int = 0
2792     same_count: int = 0
2793     failure_count: int = 0
2794
2795     def done(self, src: Path, changed: Changed) -> None:
2796         """Increment the counter for successful reformatting. Write out a message."""
2797         if changed is Changed.YES:
2798             reformatted = "would reformat" if self.check else "reformatted"
2799             if not self.quiet:
2800                 out(f"{reformatted} {src}")
2801             self.change_count += 1
2802         else:
2803             if not self.quiet:
2804                 if changed is Changed.NO:
2805                     msg = f"{src} already well formatted, good job."
2806                 else:
2807                     msg = f"{src} wasn't modified on disk since last run."
2808                 out(msg, bold=False)
2809             self.same_count += 1
2810
2811     def failed(self, src: Path, message: str) -> None:
2812         """Increment the counter for failed reformatting. Write out a message."""
2813         err(f"error: cannot format {src}: {message}")
2814         self.failure_count += 1
2815
2816     @property
2817     def return_code(self) -> int:
2818         """Return the exit code that the app should use.
2819
2820         This considers the current state of changed files and failures:
2821         - if there were any failures, return 123;
2822         - if any files were changed and --check is being used, return 1;
2823         - otherwise return 0.
2824         """
2825         # According to http://tldp.org/LDP/abs/html/exitcodes.html starting with
2826         # 126 we have special returncodes reserved by the shell.
2827         if self.failure_count:
2828             return 123
2829
2830         elif self.change_count and self.check:
2831             return 1
2832
2833         return 0
2834
2835     def __str__(self) -> str:
2836         """Render a color report of the current state.
2837
2838         Use `click.unstyle` to remove colors.
2839         """
2840         if self.check:
2841             reformatted = "would be reformatted"
2842             unchanged = "would be left unchanged"
2843             failed = "would fail to reformat"
2844         else:
2845             reformatted = "reformatted"
2846             unchanged = "left unchanged"
2847             failed = "failed to reformat"
2848         report = []
2849         if self.change_count:
2850             s = "s" if self.change_count > 1 else ""
2851             report.append(
2852                 click.style(f"{self.change_count} file{s} {reformatted}", bold=True)
2853             )
2854         if self.same_count:
2855             s = "s" if self.same_count > 1 else ""
2856             report.append(f"{self.same_count} file{s} {unchanged}")
2857         if self.failure_count:
2858             s = "s" if self.failure_count > 1 else ""
2859             report.append(
2860                 click.style(f"{self.failure_count} file{s} {failed}", fg="red")
2861             )
2862         return ", ".join(report) + "."
2863
2864
2865 def assert_equivalent(src: str, dst: str) -> None:
2866     """Raise AssertionError if `src` and `dst` aren't equivalent."""
2867
2868     import ast
2869     import traceback
2870
2871     def _v(node: ast.AST, depth: int = 0) -> Iterator[str]:
2872         """Simple visitor generating strings to compare ASTs by content."""
2873         yield f"{'  ' * depth}{node.__class__.__name__}("
2874
2875         for field in sorted(node._fields):
2876             try:
2877                 value = getattr(node, field)
2878             except AttributeError:
2879                 continue
2880
2881             yield f"{'  ' * (depth+1)}{field}="
2882
2883             if isinstance(value, list):
2884                 for item in value:
2885                     if isinstance(item, ast.AST):
2886                         yield from _v(item, depth + 2)
2887
2888             elif isinstance(value, ast.AST):
2889                 yield from _v(value, depth + 2)
2890
2891             else:
2892                 yield f"{'  ' * (depth+2)}{value!r},  # {value.__class__.__name__}"
2893
2894         yield f"{'  ' * depth})  # /{node.__class__.__name__}"
2895
2896     try:
2897         src_ast = ast.parse(src)
2898     except Exception as exc:
2899         major, minor = sys.version_info[:2]
2900         raise AssertionError(
2901             f"cannot use --safe with this file; failed to parse source file "
2902             f"with Python {major}.{minor}'s builtin AST. Re-run with --fast "
2903             f"or stop using deprecated Python 2 syntax. AST error message: {exc}"
2904         )
2905
2906     try:
2907         dst_ast = ast.parse(dst)
2908     except Exception as exc:
2909         log = dump_to_file("".join(traceback.format_tb(exc.__traceback__)), dst)
2910         raise AssertionError(
2911             f"INTERNAL ERROR: Black produced invalid code: {exc}. "
2912             f"Please report a bug on https://github.com/ambv/black/issues.  "
2913             f"This invalid output might be helpful: {log}"
2914         ) from None
2915
2916     src_ast_str = "\n".join(_v(src_ast))
2917     dst_ast_str = "\n".join(_v(dst_ast))
2918     if src_ast_str != dst_ast_str:
2919         log = dump_to_file(diff(src_ast_str, dst_ast_str, "src", "dst"))
2920         raise AssertionError(
2921             f"INTERNAL ERROR: Black produced code that is not equivalent to "
2922             f"the source.  "
2923             f"Please report a bug on https://github.com/ambv/black/issues.  "
2924             f"This diff might be helpful: {log}"
2925         ) from None
2926
2927
2928 def assert_stable(
2929     src: str, dst: str, line_length: int, is_pyi: bool = False, force_py36: bool = False
2930 ) -> None:
2931     """Raise AssertionError if `dst` reformats differently the second time."""
2932     newdst = format_str(
2933         dst, line_length=line_length, is_pyi=is_pyi, force_py36=force_py36
2934     )
2935     if dst != newdst:
2936         log = dump_to_file(
2937             diff(src, dst, "source", "first pass"),
2938             diff(dst, newdst, "first pass", "second pass"),
2939         )
2940         raise AssertionError(
2941             f"INTERNAL ERROR: Black produced different code on the second pass "
2942             f"of the formatter.  "
2943             f"Please report a bug on https://github.com/ambv/black/issues.  "
2944             f"This diff might be helpful: {log}"
2945         ) from None
2946
2947
2948 def dump_to_file(*output: str) -> str:
2949     """Dump `output` to a temporary file. Return path to the file."""
2950     import tempfile
2951
2952     with tempfile.NamedTemporaryFile(
2953         mode="w", prefix="blk_", suffix=".log", delete=False, encoding="utf8"
2954     ) as f:
2955         for lines in output:
2956             f.write(lines)
2957             if lines and lines[-1] != "\n":
2958                 f.write("\n")
2959     return f.name
2960
2961
2962 def diff(a: str, b: str, a_name: str, b_name: str) -> str:
2963     """Return a unified diff string between strings `a` and `b`."""
2964     import difflib
2965
2966     a_lines = [line + "\n" for line in a.split("\n")]
2967     b_lines = [line + "\n" for line in b.split("\n")]
2968     return "".join(
2969         difflib.unified_diff(a_lines, b_lines, fromfile=a_name, tofile=b_name, n=5)
2970     )
2971
2972
2973 def cancel(tasks: Iterable[asyncio.Task]) -> None:
2974     """asyncio signal handler that cancels all `tasks` and reports to stderr."""
2975     err("Aborted!")
2976     for task in tasks:
2977         task.cancel()
2978
2979
2980 def shutdown(loop: BaseEventLoop) -> None:
2981     """Cancel all pending tasks on `loop`, wait for them, and close the loop."""
2982     try:
2983         # This part is borrowed from asyncio/runners.py in Python 3.7b2.
2984         to_cancel = [task for task in asyncio.Task.all_tasks(loop) if not task.done()]
2985         if not to_cancel:
2986             return
2987
2988         for task in to_cancel:
2989             task.cancel()
2990         loop.run_until_complete(
2991             asyncio.gather(*to_cancel, loop=loop, return_exceptions=True)
2992         )
2993     finally:
2994         # `concurrent.futures.Future` objects cannot be cancelled once they
2995         # are already running. There might be some when the `shutdown()` happened.
2996         # Silence their logger's spew about the event loop being closed.
2997         cf_logger = logging.getLogger("concurrent.futures")
2998         cf_logger.setLevel(logging.CRITICAL)
2999         loop.close()
3000
3001
3002 def sub_twice(regex: Pattern[str], replacement: str, original: str) -> str:
3003     """Replace `regex` with `replacement` twice on `original`.
3004
3005     This is used by string normalization to perform replaces on
3006     overlapping matches.
3007     """
3008     return regex.sub(replacement, regex.sub(replacement, original))
3009
3010
3011 def enumerate_reversed(sequence: Sequence[T]) -> Iterator[Tuple[Index, T]]:
3012     """Like `reversed(enumerate(sequence))` if that were possible."""
3013     index = len(sequence) - 1
3014     for element in reversed(sequence):
3015         yield (index, element)
3016         index -= 1
3017
3018
3019 def enumerate_with_length(
3020     line: Line, reversed: bool = False
3021 ) -> Iterator[Tuple[Index, Leaf, int]]:
3022     """Return an enumeration of leaves with their length.
3023
3024     Stops prematurely on multiline strings and standalone comments.
3025     """
3026     op = cast(
3027         Callable[[Sequence[Leaf]], Iterator[Tuple[Index, Leaf]]],
3028         enumerate_reversed if reversed else enumerate,
3029     )
3030     for index, leaf in op(line.leaves):
3031         length = len(leaf.prefix) + len(leaf.value)
3032         if "\n" in leaf.value:
3033             return  # Multiline strings, we can't continue.
3034
3035         comment: Optional[Leaf]
3036         for comment in line.comments_after(leaf, index):
3037             length += len(comment.value)
3038
3039         yield index, leaf, length
3040
3041
3042 def is_line_short_enough(line: Line, *, line_length: int, line_str: str = "") -> bool:
3043     """Return True if `line` is no longer than `line_length`.
3044
3045     Uses the provided `line_str` rendering, if any, otherwise computes a new one.
3046     """
3047     if not line_str:
3048         line_str = str(line).strip("\n")
3049     return (
3050         len(line_str) <= line_length
3051         and "\n" not in line_str  # multiline strings
3052         and not line.contains_standalone_comments()
3053     )
3054
3055
3056 def can_omit_invisible_parens(line: Line, line_length: int) -> bool:
3057     """Does `line` have a shape safe to reformat without optional parens around it?
3058
3059     Returns True for only a subset of potentially nice looking formattings but
3060     the point is to not return false positives that end up producing lines that
3061     are too long.
3062     """
3063     bt = line.bracket_tracker
3064     if not bt.delimiters:
3065         # Without delimiters the optional parentheses are useless.
3066         return True
3067
3068     max_priority = bt.max_delimiter_priority()
3069     if bt.delimiter_count_with_priority(max_priority) > 1:
3070         # With more than one delimiter of a kind the optional parentheses read better.
3071         return False
3072
3073     if max_priority == DOT_PRIORITY:
3074         # A single stranded method call doesn't require optional parentheses.
3075         return True
3076
3077     assert len(line.leaves) >= 2, "Stranded delimiter"
3078
3079     first = line.leaves[0]
3080     second = line.leaves[1]
3081     penultimate = line.leaves[-2]
3082     last = line.leaves[-1]
3083
3084     # With a single delimiter, omit if the expression starts or ends with
3085     # a bracket.
3086     if first.type in OPENING_BRACKETS and second.type not in CLOSING_BRACKETS:
3087         remainder = False
3088         length = 4 * line.depth
3089         for _index, leaf, leaf_length in enumerate_with_length(line):
3090             if leaf.type in CLOSING_BRACKETS and leaf.opening_bracket is first:
3091                 remainder = True
3092             if remainder:
3093                 length += leaf_length
3094                 if length > line_length:
3095                     break
3096
3097                 if leaf.type in OPENING_BRACKETS:
3098                     # There are brackets we can further split on.
3099                     remainder = False
3100
3101         else:
3102             # checked the entire string and line length wasn't exceeded
3103             if len(line.leaves) == _index + 1:
3104                 return True
3105
3106         # Note: we are not returning False here because a line might have *both*
3107         # a leading opening bracket and a trailing closing bracket.  If the
3108         # opening bracket doesn't match our rule, maybe the closing will.
3109
3110     if (
3111         last.type == token.RPAR
3112         or last.type == token.RBRACE
3113         or (
3114             # don't use indexing for omitting optional parentheses;
3115             # it looks weird
3116             last.type == token.RSQB
3117             and last.parent
3118             and last.parent.type != syms.trailer
3119         )
3120     ):
3121         if penultimate.type in OPENING_BRACKETS:
3122             # Empty brackets don't help.
3123             return False
3124
3125         if is_multiline_string(first):
3126             # Additional wrapping of a multiline string in this situation is
3127             # unnecessary.
3128             return True
3129
3130         length = 4 * line.depth
3131         seen_other_brackets = False
3132         for _index, leaf, leaf_length in enumerate_with_length(line):
3133             length += leaf_length
3134             if leaf is last.opening_bracket:
3135                 if seen_other_brackets or length <= line_length:
3136                     return True
3137
3138             elif leaf.type in OPENING_BRACKETS:
3139                 # There are brackets we can further split on.
3140                 seen_other_brackets = True
3141
3142     return False
3143
3144
3145 def get_cache_file(line_length: int, pyi: bool = False, py36: bool = False) -> Path:
3146     return (
3147         CACHE_DIR
3148         / f"cache.{line_length}{'.pyi' if pyi else ''}{'.py36' if py36 else ''}.pickle"
3149     )
3150
3151
3152 def read_cache(line_length: int, pyi: bool = False, py36: bool = False) -> Cache:
3153     """Read the cache if it exists and is well formed.
3154
3155     If it is not well formed, the call to write_cache later should resolve the issue.
3156     """
3157     cache_file = get_cache_file(line_length, pyi, py36)
3158     if not cache_file.exists():
3159         return {}
3160
3161     with cache_file.open("rb") as fobj:
3162         try:
3163             cache: Cache = pickle.load(fobj)
3164         except pickle.UnpicklingError:
3165             return {}
3166
3167     return cache
3168
3169
3170 def get_cache_info(path: Path) -> CacheInfo:
3171     """Return the information used to check if a file is already formatted or not."""
3172     stat = path.stat()
3173     return stat.st_mtime, stat.st_size
3174
3175
3176 def filter_cached(
3177     cache: Cache, sources: Iterable[Path]
3178 ) -> Tuple[List[Path], List[Path]]:
3179     """Split a list of paths into two.
3180
3181     The first list contains paths of files that modified on disk or are not in the
3182     cache. The other list contains paths to non-modified files.
3183     """
3184     todo, done = [], []
3185     for src in sources:
3186         src = src.resolve()
3187         if cache.get(src) != get_cache_info(src):
3188             todo.append(src)
3189         else:
3190             done.append(src)
3191     return todo, done
3192
3193
3194 def write_cache(
3195     cache: Cache,
3196     sources: List[Path],
3197     line_length: int,
3198     pyi: bool = False,
3199     py36: bool = False,
3200 ) -> None:
3201     """Update the cache file."""
3202     cache_file = get_cache_file(line_length, pyi, py36)
3203     try:
3204         if not CACHE_DIR.exists():
3205             CACHE_DIR.mkdir(parents=True)
3206         new_cache = {**cache, **{src.resolve(): get_cache_info(src) for src in sources}}
3207         with cache_file.open("wb") as fobj:
3208             pickle.dump(new_cache, fobj, protocol=pickle.HIGHEST_PROTOCOL)
3209     except OSError:
3210         pass
3211
3212
3213 if __name__ == "__main__":
3214     main()