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

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