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

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