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

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