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

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