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

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