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

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