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

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