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

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