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

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