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

c74ba4e0e8465573150450b7cc09e985f7bb12d6
[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.6b1"
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 = (
1260             leaf if leaf.type == token.LSQB else self.bracket_tracker.get_open_lsqb()
1261         )
1262         if open_lsqb is None:
1263             return False
1264
1265         subscript_start = open_lsqb.next_sibling
1266         if (
1267             isinstance(subscript_start, Node)
1268             and subscript_start.type == syms.subscriptlist
1269         ):
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 quotes
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 quotes 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 new_quote == '"""' and new_body[-1:] == '"':
2553         # edge case:
2554         new_body = new_body[:-1] + '\\"'
2555     orig_escape_count = body.count("\\")
2556     new_escape_count = new_body.count("\\")
2557     if new_escape_count > orig_escape_count:
2558         return  # Do not introduce more escaping
2559
2560     if new_escape_count == orig_escape_count and orig_quote == '"':
2561         return  # Prefer double quotes
2562
2563     leaf.value = f"{prefix}{new_quote}{new_body}{new_quote}"
2564
2565
2566 def normalize_invisible_parens(node: Node, parens_after: Set[str]) -> None:
2567     """Make existing optional parentheses invisible or create new ones.
2568
2569     `parens_after` is a set of string leaf values immeditely after which parens
2570     should be put.
2571
2572     Standardizes on visible parentheses for single-element tuples, and keeps
2573     existing visible parentheses for other tuples and generator expressions.
2574     """
2575     try:
2576         list(generate_comments(node))
2577     except FormatOff:
2578         return  # This `node` has a prefix with `# fmt: off`, don't mess with parens.
2579
2580     check_lpar = False
2581     for index, child in enumerate(list(node.children)):
2582         if check_lpar:
2583             if child.type == syms.atom:
2584                 maybe_make_parens_invisible_in_atom(child)
2585             elif is_one_tuple(child):
2586                 # wrap child in visible parentheses
2587                 lpar = Leaf(token.LPAR, "(")
2588                 rpar = Leaf(token.RPAR, ")")
2589                 child.remove()
2590                 node.insert_child(index, Node(syms.atom, [lpar, child, rpar]))
2591             elif node.type == syms.import_from:
2592                 # "import from" nodes store parentheses directly as part of
2593                 # the statement
2594                 if child.type == token.LPAR:
2595                     # make parentheses invisible
2596                     child.value = ""  # type: ignore
2597                     node.children[-1].value = ""  # type: ignore
2598                 elif child.type != token.STAR:
2599                     # insert invisible parentheses
2600                     node.insert_child(index, Leaf(token.LPAR, ""))
2601                     node.append_child(Leaf(token.RPAR, ""))
2602                 break
2603
2604             elif not (isinstance(child, Leaf) and is_multiline_string(child)):
2605                 # wrap child in invisible parentheses
2606                 lpar = Leaf(token.LPAR, "")
2607                 rpar = Leaf(token.RPAR, "")
2608                 index = child.remove() or 0
2609                 node.insert_child(index, Node(syms.atom, [lpar, child, rpar]))
2610
2611         check_lpar = isinstance(child, Leaf) and child.value in parens_after
2612
2613
2614 def maybe_make_parens_invisible_in_atom(node: LN) -> bool:
2615     """If it's safe, make the parens in the atom `node` invisible, recursively."""
2616     if (
2617         node.type != syms.atom
2618         or is_empty_tuple(node)
2619         or is_one_tuple(node)
2620         or is_yield(node)
2621         or max_delimiter_priority_in_atom(node) >= COMMA_PRIORITY
2622     ):
2623         return False
2624
2625     first = node.children[0]
2626     last = node.children[-1]
2627     if first.type == token.LPAR and last.type == token.RPAR:
2628         # make parentheses invisible
2629         first.value = ""  # type: ignore
2630         last.value = ""  # type: ignore
2631         if len(node.children) > 1:
2632             maybe_make_parens_invisible_in_atom(node.children[1])
2633         return True
2634
2635     return False
2636
2637
2638 def is_empty_tuple(node: LN) -> bool:
2639     """Return True if `node` holds an empty tuple."""
2640     return (
2641         node.type == syms.atom
2642         and len(node.children) == 2
2643         and node.children[0].type == token.LPAR
2644         and node.children[1].type == token.RPAR
2645     )
2646
2647
2648 def is_one_tuple(node: LN) -> bool:
2649     """Return True if `node` holds a tuple with one element, with or without parens."""
2650     if node.type == syms.atom:
2651         if len(node.children) != 3:
2652             return False
2653
2654         lpar, gexp, rpar = node.children
2655         if not (
2656             lpar.type == token.LPAR
2657             and gexp.type == syms.testlist_gexp
2658             and rpar.type == token.RPAR
2659         ):
2660             return False
2661
2662         return len(gexp.children) == 2 and gexp.children[1].type == token.COMMA
2663
2664     return (
2665         node.type in IMPLICIT_TUPLE
2666         and len(node.children) == 2
2667         and node.children[1].type == token.COMMA
2668     )
2669
2670
2671 def is_yield(node: LN) -> bool:
2672     """Return True if `node` holds a `yield` or `yield from` expression."""
2673     if node.type == syms.yield_expr:
2674         return True
2675
2676     if node.type == token.NAME and node.value == "yield":  # type: ignore
2677         return True
2678
2679     if node.type != syms.atom:
2680         return False
2681
2682     if len(node.children) != 3:
2683         return False
2684
2685     lpar, expr, rpar = node.children
2686     if lpar.type == token.LPAR and rpar.type == token.RPAR:
2687         return is_yield(expr)
2688
2689     return False
2690
2691
2692 def is_vararg(leaf: Leaf, within: Set[NodeType]) -> bool:
2693     """Return True if `leaf` is a star or double star in a vararg or kwarg.
2694
2695     If `within` includes VARARGS_PARENTS, this applies to function signatures.
2696     If `within` includes UNPACKING_PARENTS, it applies to right hand-side
2697     extended iterable unpacking (PEP 3132) and additional unpacking
2698     generalizations (PEP 448).
2699     """
2700     if leaf.type not in STARS or not leaf.parent:
2701         return False
2702
2703     p = leaf.parent
2704     if p.type == syms.star_expr:
2705         # Star expressions are also used as assignment targets in extended
2706         # iterable unpacking (PEP 3132).  See what its parent is instead.
2707         if not p.parent:
2708             return False
2709
2710         p = p.parent
2711
2712     return p.type in within
2713
2714
2715 def is_multiline_string(leaf: Leaf) -> bool:
2716     """Return True if `leaf` is a multiline string that actually spans many lines."""
2717     value = leaf.value.lstrip("furbFURB")
2718     return value[:3] in {'"""', "'''"} and "\n" in value
2719
2720
2721 def is_stub_suite(node: Node) -> bool:
2722     """Return True if `node` is a suite with a stub body."""
2723     if (
2724         len(node.children) != 4
2725         or node.children[0].type != token.NEWLINE
2726         or node.children[1].type != token.INDENT
2727         or node.children[3].type != token.DEDENT
2728     ):
2729         return False
2730
2731     return is_stub_body(node.children[2])
2732
2733
2734 def is_stub_body(node: LN) -> bool:
2735     """Return True if `node` is a simple statement containing an ellipsis."""
2736     if not isinstance(node, Node) or node.type != syms.simple_stmt:
2737         return False
2738
2739     if len(node.children) != 2:
2740         return False
2741
2742     child = node.children[0]
2743     return (
2744         child.type == syms.atom
2745         and len(child.children) == 3
2746         and all(leaf == Leaf(token.DOT, ".") for leaf in child.children)
2747     )
2748
2749
2750 def max_delimiter_priority_in_atom(node: LN) -> int:
2751     """Return maximum delimiter priority inside `node`.
2752
2753     This is specific to atoms with contents contained in a pair of parentheses.
2754     If `node` isn't an atom or there are no enclosing parentheses, returns 0.
2755     """
2756     if node.type != syms.atom:
2757         return 0
2758
2759     first = node.children[0]
2760     last = node.children[-1]
2761     if not (first.type == token.LPAR and last.type == token.RPAR):
2762         return 0
2763
2764     bt = BracketTracker()
2765     for c in node.children[1:-1]:
2766         if isinstance(c, Leaf):
2767             bt.mark(c)
2768         else:
2769             for leaf in c.leaves():
2770                 bt.mark(leaf)
2771     try:
2772         return bt.max_delimiter_priority()
2773
2774     except ValueError:
2775         return 0
2776
2777
2778 def ensure_visible(leaf: Leaf) -> None:
2779     """Make sure parentheses are visible.
2780
2781     They could be invisible as part of some statements (see
2782     :func:`normalize_invible_parens` and :func:`visit_import_from`).
2783     """
2784     if leaf.type == token.LPAR:
2785         leaf.value = "("
2786     elif leaf.type == token.RPAR:
2787         leaf.value = ")"
2788
2789
2790 def should_explode(line: Line, opening_bracket: Leaf) -> bool:
2791     """Should `line` immediately be split with `delimiter_split()` after RHS?"""
2792     if not (
2793         opening_bracket.parent
2794         and opening_bracket.parent.type in {syms.atom, syms.import_from}
2795         and opening_bracket.value in "[{("
2796     ):
2797         return False
2798
2799     try:
2800         last_leaf = line.leaves[-1]
2801         exclude = {id(last_leaf)} if last_leaf.type == token.COMMA else set()
2802         max_priority = line.bracket_tracker.max_delimiter_priority(exclude=exclude)
2803     except (IndexError, ValueError):
2804         return False
2805
2806     return max_priority == COMMA_PRIORITY
2807
2808
2809 def is_python36(node: Node) -> bool:
2810     """Return True if the current file is using Python 3.6+ features.
2811
2812     Currently looking for:
2813     - f-strings; and
2814     - trailing commas after * or ** in function signatures and calls.
2815     """
2816     for n in node.pre_order():
2817         if n.type == token.STRING:
2818             value_head = n.value[:2]  # type: ignore
2819             if value_head in {'f"', 'F"', "f'", "F'", "rf", "fr", "RF", "FR"}:
2820                 return True
2821
2822         elif (
2823             n.type in {syms.typedargslist, syms.arglist}
2824             and n.children
2825             and n.children[-1].type == token.COMMA
2826         ):
2827             for ch in n.children:
2828                 if ch.type in STARS:
2829                     return True
2830
2831                 if ch.type == syms.argument:
2832                     for argch in ch.children:
2833                         if argch.type in STARS:
2834                             return True
2835
2836     return False
2837
2838
2839 def generate_trailers_to_omit(line: Line, line_length: int) -> Iterator[Set[LeafID]]:
2840     """Generate sets of closing bracket IDs that should be omitted in a RHS.
2841
2842     Brackets can be omitted if the entire trailer up to and including
2843     a preceding closing bracket fits in one line.
2844
2845     Yielded sets are cumulative (contain results of previous yields, too).  First
2846     set is empty.
2847     """
2848
2849     omit: Set[LeafID] = set()
2850     yield omit
2851
2852     length = 4 * line.depth
2853     opening_bracket = None
2854     closing_bracket = None
2855     optional_brackets: Set[LeafID] = set()
2856     inner_brackets: Set[LeafID] = set()
2857     for index, leaf, leaf_length in enumerate_with_length(line, reversed=True):
2858         length += leaf_length
2859         if length > line_length:
2860             break
2861
2862         has_inline_comment = leaf_length > len(leaf.value) + len(leaf.prefix)
2863         if leaf.type == STANDALONE_COMMENT or has_inline_comment:
2864             break
2865
2866         optional_brackets.discard(id(leaf))
2867         if opening_bracket:
2868             if leaf is opening_bracket:
2869                 opening_bracket = None
2870             elif leaf.type in CLOSING_BRACKETS:
2871                 inner_brackets.add(id(leaf))
2872         elif leaf.type in CLOSING_BRACKETS:
2873             if not leaf.value:
2874                 optional_brackets.add(id(opening_bracket))
2875                 continue
2876
2877             if index > 0 and line.leaves[index - 1].type in OPENING_BRACKETS:
2878                 # Empty brackets would fail a split so treat them as "inner"
2879                 # brackets (e.g. only add them to the `omit` set if another
2880                 # pair of brackets was good enough.
2881                 inner_brackets.add(id(leaf))
2882                 continue
2883
2884             opening_bracket = leaf.opening_bracket
2885             if closing_bracket:
2886                 omit.add(id(closing_bracket))
2887                 omit.update(inner_brackets)
2888                 inner_brackets.clear()
2889                 yield omit
2890             closing_bracket = leaf
2891
2892
2893 def get_future_imports(node: Node) -> Set[str]:
2894     """Return a set of __future__ imports in the file."""
2895     imports = set()
2896     for child in node.children:
2897         if child.type != syms.simple_stmt:
2898             break
2899         first_child = child.children[0]
2900         if isinstance(first_child, Leaf):
2901             # Continue looking if we see a docstring; otherwise stop.
2902             if (
2903                 len(child.children) == 2
2904                 and first_child.type == token.STRING
2905                 and child.children[1].type == token.NEWLINE
2906             ):
2907                 continue
2908             else:
2909                 break
2910         elif first_child.type == syms.import_from:
2911             module_name = first_child.children[1]
2912             if not isinstance(module_name, Leaf) or module_name.value != "__future__":
2913                 break
2914             for import_from_child in first_child.children[3:]:
2915                 if isinstance(import_from_child, Leaf):
2916                     if import_from_child.type == token.NAME:
2917                         imports.add(import_from_child.value)
2918                 else:
2919                     assert import_from_child.type == syms.import_as_names
2920                     for leaf in import_from_child.children:
2921                         if isinstance(leaf, Leaf) and leaf.type == token.NAME:
2922                             imports.add(leaf.value)
2923         else:
2924             break
2925     return imports
2926
2927
2928 def gen_python_files_in_dir(
2929     path: Path,
2930     root: Path,
2931     include: Pattern[str],
2932     exclude: Pattern[str],
2933     report: "Report",
2934 ) -> Iterator[Path]:
2935     """Generate all files under `path` whose paths are not excluded by the
2936     `exclude` regex, but are included by the `include` regex.
2937
2938     `report` is where output about exclusions goes.
2939     """
2940     assert root.is_absolute(), f"INTERNAL ERROR: `root` must be absolute but is {root}"
2941     for child in path.iterdir():
2942         normalized_path = "/" + child.resolve().relative_to(root).as_posix()
2943         if child.is_dir():
2944             normalized_path += "/"
2945         exclude_match = exclude.search(normalized_path)
2946         if exclude_match and exclude_match.group(0):
2947             report.path_ignored(child, f"matches the --exclude regular expression")
2948             continue
2949
2950         if child.is_dir():
2951             yield from gen_python_files_in_dir(child, root, include, exclude, report)
2952
2953         elif child.is_file():
2954             include_match = include.search(normalized_path)
2955             if include_match:
2956                 yield child
2957
2958
2959 @lru_cache()
2960 def find_project_root(srcs: Iterable[str]) -> Path:
2961     """Return a directory containing .git, .hg, or pyproject.toml.
2962
2963     That directory can be one of the directories passed in `srcs` or their
2964     common parent.
2965
2966     If no directory in the tree contains a marker that would specify it's the
2967     project root, the root of the file system is returned.
2968     """
2969     if not srcs:
2970         return Path("/").resolve()
2971
2972     common_base = min(Path(src).resolve() for src in srcs)
2973     if common_base.is_dir():
2974         # Append a fake file so `parents` below returns `common_base_dir`, too.
2975         common_base /= "fake-file"
2976     for directory in common_base.parents:
2977         if (directory / ".git").is_dir():
2978             return directory
2979
2980         if (directory / ".hg").is_dir():
2981             return directory
2982
2983         if (directory / "pyproject.toml").is_file():
2984             return directory
2985
2986     return directory
2987
2988
2989 @dataclass
2990 class Report:
2991     """Provides a reformatting counter. Can be rendered with `str(report)`."""
2992
2993     check: bool = False
2994     quiet: bool = False
2995     verbose: bool = False
2996     change_count: int = 0
2997     same_count: int = 0
2998     failure_count: int = 0
2999
3000     def done(self, src: Path, changed: Changed) -> None:
3001         """Increment the counter for successful reformatting. Write out a message."""
3002         if changed is Changed.YES:
3003             reformatted = "would reformat" if self.check else "reformatted"
3004             if self.verbose or not self.quiet:
3005                 out(f"{reformatted} {src}")
3006             self.change_count += 1
3007         else:
3008             if self.verbose:
3009                 if changed is Changed.NO:
3010                     msg = f"{src} already well formatted, good job."
3011                 else:
3012                     msg = f"{src} wasn't modified on disk since last run."
3013                 out(msg, bold=False)
3014             self.same_count += 1
3015
3016     def failed(self, src: Path, message: str) -> None:
3017         """Increment the counter for failed reformatting. Write out a message."""
3018         err(f"error: cannot format {src}: {message}")
3019         self.failure_count += 1
3020
3021     def path_ignored(self, path: Path, message: str) -> None:
3022         if self.verbose:
3023             out(f"{path} ignored: {message}", bold=False)
3024
3025     @property
3026     def return_code(self) -> int:
3027         """Return the exit code that the app should use.
3028
3029         This considers the current state of changed files and failures:
3030         - if there were any failures, return 123;
3031         - if any files were changed and --check is being used, return 1;
3032         - otherwise return 0.
3033         """
3034         # According to http://tldp.org/LDP/abs/html/exitcodes.html starting with
3035         # 126 we have special returncodes reserved by the shell.
3036         if self.failure_count:
3037             return 123
3038
3039         elif self.change_count and self.check:
3040             return 1
3041
3042         return 0
3043
3044     def __str__(self) -> str:
3045         """Render a color report of the current state.
3046
3047         Use `click.unstyle` to remove colors.
3048         """
3049         if self.check:
3050             reformatted = "would be reformatted"
3051             unchanged = "would be left unchanged"
3052             failed = "would fail to reformat"
3053         else:
3054             reformatted = "reformatted"
3055             unchanged = "left unchanged"
3056             failed = "failed to reformat"
3057         report = []
3058         if self.change_count:
3059             s = "s" if self.change_count > 1 else ""
3060             report.append(
3061                 click.style(f"{self.change_count} file{s} {reformatted}", bold=True)
3062             )
3063         if self.same_count:
3064             s = "s" if self.same_count > 1 else ""
3065             report.append(f"{self.same_count} file{s} {unchanged}")
3066         if self.failure_count:
3067             s = "s" if self.failure_count > 1 else ""
3068             report.append(
3069                 click.style(f"{self.failure_count} file{s} {failed}", fg="red")
3070             )
3071         return ", ".join(report) + "."
3072
3073
3074 def assert_equivalent(src: str, dst: str) -> None:
3075     """Raise AssertionError if `src` and `dst` aren't equivalent."""
3076
3077     import ast
3078     import traceback
3079
3080     def _v(node: ast.AST, depth: int = 0) -> Iterator[str]:
3081         """Simple visitor generating strings to compare ASTs by content."""
3082         yield f"{'  ' * depth}{node.__class__.__name__}("
3083
3084         for field in sorted(node._fields):
3085             try:
3086                 value = getattr(node, field)
3087             except AttributeError:
3088                 continue
3089
3090             yield f"{'  ' * (depth+1)}{field}="
3091
3092             if isinstance(value, list):
3093                 for item in value:
3094                     if isinstance(item, ast.AST):
3095                         yield from _v(item, depth + 2)
3096
3097             elif isinstance(value, ast.AST):
3098                 yield from _v(value, depth + 2)
3099
3100             else:
3101                 yield f"{'  ' * (depth+2)}{value!r},  # {value.__class__.__name__}"
3102
3103         yield f"{'  ' * depth})  # /{node.__class__.__name__}"
3104
3105     try:
3106         src_ast = ast.parse(src)
3107     except Exception as exc:
3108         major, minor = sys.version_info[:2]
3109         raise AssertionError(
3110             f"cannot use --safe with this file; failed to parse source file "
3111             f"with Python {major}.{minor}'s builtin AST. Re-run with --fast "
3112             f"or stop using deprecated Python 2 syntax. AST error message: {exc}"
3113         )
3114
3115     try:
3116         dst_ast = ast.parse(dst)
3117     except Exception as exc:
3118         log = dump_to_file("".join(traceback.format_tb(exc.__traceback__)), dst)
3119         raise AssertionError(
3120             f"INTERNAL ERROR: Black produced invalid code: {exc}. "
3121             f"Please report a bug on https://github.com/ambv/black/issues.  "
3122             f"This invalid output might be helpful: {log}"
3123         ) from None
3124
3125     src_ast_str = "\n".join(_v(src_ast))
3126     dst_ast_str = "\n".join(_v(dst_ast))
3127     if src_ast_str != dst_ast_str:
3128         log = dump_to_file(diff(src_ast_str, dst_ast_str, "src", "dst"))
3129         raise AssertionError(
3130             f"INTERNAL ERROR: Black produced code that is not equivalent to "
3131             f"the source.  "
3132             f"Please report a bug on https://github.com/ambv/black/issues.  "
3133             f"This diff might be helpful: {log}"
3134         ) from None
3135
3136
3137 def assert_stable(
3138     src: str, dst: str, line_length: int, mode: FileMode = FileMode.AUTO_DETECT
3139 ) -> None:
3140     """Raise AssertionError if `dst` reformats differently the second time."""
3141     newdst = format_str(dst, line_length=line_length, mode=mode)
3142     if dst != newdst:
3143         log = dump_to_file(
3144             diff(src, dst, "source", "first pass"),
3145             diff(dst, newdst, "first pass", "second pass"),
3146         )
3147         raise AssertionError(
3148             f"INTERNAL ERROR: Black produced different code on the second pass "
3149             f"of the formatter.  "
3150             f"Please report a bug on https://github.com/ambv/black/issues.  "
3151             f"This diff might be helpful: {log}"
3152         ) from None
3153
3154
3155 def dump_to_file(*output: str) -> str:
3156     """Dump `output` to a temporary file. Return path to the file."""
3157     import tempfile
3158
3159     with tempfile.NamedTemporaryFile(
3160         mode="w", prefix="blk_", suffix=".log", delete=False, encoding="utf8"
3161     ) as f:
3162         for lines in output:
3163             f.write(lines)
3164             if lines and lines[-1] != "\n":
3165                 f.write("\n")
3166     return f.name
3167
3168
3169 def diff(a: str, b: str, a_name: str, b_name: str) -> str:
3170     """Return a unified diff string between strings `a` and `b`."""
3171     import difflib
3172
3173     a_lines = [line + "\n" for line in a.split("\n")]
3174     b_lines = [line + "\n" for line in b.split("\n")]
3175     return "".join(
3176         difflib.unified_diff(a_lines, b_lines, fromfile=a_name, tofile=b_name, n=5)
3177     )
3178
3179
3180 def cancel(tasks: Iterable[asyncio.Task]) -> None:
3181     """asyncio signal handler that cancels all `tasks` and reports to stderr."""
3182     err("Aborted!")
3183     for task in tasks:
3184         task.cancel()
3185
3186
3187 def shutdown(loop: BaseEventLoop) -> None:
3188     """Cancel all pending tasks on `loop`, wait for them, and close the loop."""
3189     try:
3190         # This part is borrowed from asyncio/runners.py in Python 3.7b2.
3191         to_cancel = [task for task in asyncio.Task.all_tasks(loop) if not task.done()]
3192         if not to_cancel:
3193             return
3194
3195         for task in to_cancel:
3196             task.cancel()
3197         loop.run_until_complete(
3198             asyncio.gather(*to_cancel, loop=loop, return_exceptions=True)
3199         )
3200     finally:
3201         # `concurrent.futures.Future` objects cannot be cancelled once they
3202         # are already running. There might be some when the `shutdown()` happened.
3203         # Silence their logger's spew about the event loop being closed.
3204         cf_logger = logging.getLogger("concurrent.futures")
3205         cf_logger.setLevel(logging.CRITICAL)
3206         loop.close()
3207
3208
3209 def sub_twice(regex: Pattern[str], replacement: str, original: str) -> str:
3210     """Replace `regex` with `replacement` twice on `original`.
3211
3212     This is used by string normalization to perform replaces on
3213     overlapping matches.
3214     """
3215     return regex.sub(replacement, regex.sub(replacement, original))
3216
3217
3218 def re_compile_maybe_verbose(regex: str) -> Pattern[str]:
3219     """Compile a regular expression string in `regex`.
3220
3221     If it contains newlines, use verbose mode.
3222     """
3223     if "\n" in regex:
3224         regex = "(?x)" + regex
3225     return re.compile(regex)
3226
3227
3228 def enumerate_reversed(sequence: Sequence[T]) -> Iterator[Tuple[Index, T]]:
3229     """Like `reversed(enumerate(sequence))` if that were possible."""
3230     index = len(sequence) - 1
3231     for element in reversed(sequence):
3232         yield (index, element)
3233         index -= 1
3234
3235
3236 def enumerate_with_length(
3237     line: Line, reversed: bool = False
3238 ) -> Iterator[Tuple[Index, Leaf, int]]:
3239     """Return an enumeration of leaves with their length.
3240
3241     Stops prematurely on multiline strings and standalone comments.
3242     """
3243     op = cast(
3244         Callable[[Sequence[Leaf]], Iterator[Tuple[Index, Leaf]]],
3245         enumerate_reversed if reversed else enumerate,
3246     )
3247     for index, leaf in op(line.leaves):
3248         length = len(leaf.prefix) + len(leaf.value)
3249         if "\n" in leaf.value:
3250             return  # Multiline strings, we can't continue.
3251
3252         comment: Optional[Leaf]
3253         for comment in line.comments_after(leaf, index):
3254             length += len(comment.value)
3255
3256         yield index, leaf, length
3257
3258
3259 def is_line_short_enough(line: Line, *, line_length: int, line_str: str = "") -> bool:
3260     """Return True if `line` is no longer than `line_length`.
3261
3262     Uses the provided `line_str` rendering, if any, otherwise computes a new one.
3263     """
3264     if not line_str:
3265         line_str = str(line).strip("\n")
3266     return (
3267         len(line_str) <= line_length
3268         and "\n" not in line_str  # multiline strings
3269         and not line.contains_standalone_comments()
3270     )
3271
3272
3273 def can_be_split(line: Line) -> bool:
3274     """Return False if the line cannot be split *for sure*.
3275
3276     This is not an exhaustive search but a cheap heuristic that we can use to
3277     avoid some unfortunate formattings (mostly around wrapping unsplittable code
3278     in unnecessary parentheses).
3279     """
3280     leaves = line.leaves
3281     if len(leaves) < 2:
3282         return False
3283
3284     if leaves[0].type == token.STRING and leaves[1].type == token.DOT:
3285         call_count = 0
3286         dot_count = 0
3287         next = leaves[-1]
3288         for leaf in leaves[-2::-1]:
3289             if leaf.type in OPENING_BRACKETS:
3290                 if next.type not in CLOSING_BRACKETS:
3291                     return False
3292
3293                 call_count += 1
3294             elif leaf.type == token.DOT:
3295                 dot_count += 1
3296             elif leaf.type == token.NAME:
3297                 if not (next.type == token.DOT or next.type in OPENING_BRACKETS):
3298                     return False
3299
3300             elif leaf.type not in CLOSING_BRACKETS:
3301                 return False
3302
3303             if dot_count > 1 and call_count > 1:
3304                 return False
3305
3306     return True
3307
3308
3309 def can_omit_invisible_parens(line: Line, line_length: int) -> bool:
3310     """Does `line` have a shape safe to reformat without optional parens around it?
3311
3312     Returns True for only a subset of potentially nice looking formattings but
3313     the point is to not return false positives that end up producing lines that
3314     are too long.
3315     """
3316     bt = line.bracket_tracker
3317     if not bt.delimiters:
3318         # Without delimiters the optional parentheses are useless.
3319         return True
3320
3321     max_priority = bt.max_delimiter_priority()
3322     if bt.delimiter_count_with_priority(max_priority) > 1:
3323         # With more than one delimiter of a kind the optional parentheses read better.
3324         return False
3325
3326     if max_priority == DOT_PRIORITY:
3327         # A single stranded method call doesn't require optional parentheses.
3328         return True
3329
3330     assert len(line.leaves) >= 2, "Stranded delimiter"
3331
3332     first = line.leaves[0]
3333     second = line.leaves[1]
3334     penultimate = line.leaves[-2]
3335     last = line.leaves[-1]
3336
3337     # With a single delimiter, omit if the expression starts or ends with
3338     # a bracket.
3339     if first.type in OPENING_BRACKETS and second.type not in CLOSING_BRACKETS:
3340         remainder = False
3341         length = 4 * line.depth
3342         for _index, leaf, leaf_length in enumerate_with_length(line):
3343             if leaf.type in CLOSING_BRACKETS and leaf.opening_bracket is first:
3344                 remainder = True
3345             if remainder:
3346                 length += leaf_length
3347                 if length > line_length:
3348                     break
3349
3350                 if leaf.type in OPENING_BRACKETS:
3351                     # There are brackets we can further split on.
3352                     remainder = False
3353
3354         else:
3355             # checked the entire string and line length wasn't exceeded
3356             if len(line.leaves) == _index + 1:
3357                 return True
3358
3359         # Note: we are not returning False here because a line might have *both*
3360         # a leading opening bracket and a trailing closing bracket.  If the
3361         # opening bracket doesn't match our rule, maybe the closing will.
3362
3363     if (
3364         last.type == token.RPAR
3365         or last.type == token.RBRACE
3366         or (
3367             # don't use indexing for omitting optional parentheses;
3368             # it looks weird
3369             last.type == token.RSQB
3370             and last.parent
3371             and last.parent.type != syms.trailer
3372         )
3373     ):
3374         if penultimate.type in OPENING_BRACKETS:
3375             # Empty brackets don't help.
3376             return False
3377
3378         if is_multiline_string(first):
3379             # Additional wrapping of a multiline string in this situation is
3380             # unnecessary.
3381             return True
3382
3383         length = 4 * line.depth
3384         seen_other_brackets = False
3385         for _index, leaf, leaf_length in enumerate_with_length(line):
3386             length += leaf_length
3387             if leaf is last.opening_bracket:
3388                 if seen_other_brackets or length <= line_length:
3389                     return True
3390
3391             elif leaf.type in OPENING_BRACKETS:
3392                 # There are brackets we can further split on.
3393                 seen_other_brackets = True
3394
3395     return False
3396
3397
3398 def get_cache_file(line_length: int, mode: FileMode) -> Path:
3399     return CACHE_DIR / f"cache.{line_length}.{mode.value}.pickle"
3400
3401
3402 def read_cache(line_length: int, mode: FileMode) -> Cache:
3403     """Read the cache if it exists and is well formed.
3404
3405     If it is not well formed, the call to write_cache later should resolve the issue.
3406     """
3407     cache_file = get_cache_file(line_length, mode)
3408     if not cache_file.exists():
3409         return {}
3410
3411     with cache_file.open("rb") as fobj:
3412         try:
3413             cache: Cache = pickle.load(fobj)
3414         except pickle.UnpicklingError:
3415             return {}
3416
3417     return cache
3418
3419
3420 def get_cache_info(path: Path) -> CacheInfo:
3421     """Return the information used to check if a file is already formatted or not."""
3422     stat = path.stat()
3423     return stat.st_mtime, stat.st_size
3424
3425
3426 def filter_cached(cache: Cache, sources: Iterable[Path]) -> Tuple[Set[Path], Set[Path]]:
3427     """Split an iterable of paths in `sources` into two sets.
3428
3429     The first contains paths of files that modified on disk or are not in the
3430     cache. The other contains paths to non-modified files.
3431     """
3432     todo, done = set(), set()
3433     for src in sources:
3434         src = src.resolve()
3435         if cache.get(src) != get_cache_info(src):
3436             todo.add(src)
3437         else:
3438             done.add(src)
3439     return todo, done
3440
3441
3442 def write_cache(
3443     cache: Cache, sources: Iterable[Path], line_length: int, mode: FileMode
3444 ) -> None:
3445     """Update the cache file."""
3446     cache_file = get_cache_file(line_length, mode)
3447     try:
3448         if not CACHE_DIR.exists():
3449             CACHE_DIR.mkdir(parents=True)
3450         new_cache = {**cache, **{src.resolve(): get_cache_info(src) for src in sources}}
3451         with cache_file.open("wb") as fobj:
3452             pickle.dump(new_cache, fobj, protocol=pickle.HIGHEST_PROTOCOL)
3453     except OSError:
3454         pass
3455
3456
3457 if __name__ == "__main__":
3458     main()