]> git.madduck.net Git - etc/vim.git/blob - src/black/__init__.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:

bfb77126218c303426dccf1bf8c2200cb97bbd3b
[etc/vim.git] / src / black / __init__.py
1 import ast
2 import asyncio
3 from abc import ABC, abstractmethod
4 from collections import defaultdict
5 from concurrent.futures import Executor, ThreadPoolExecutor, ProcessPoolExecutor
6 from contextlib import contextmanager
7 from datetime import datetime
8 from enum import Enum
9 from functools import lru_cache, partial, wraps
10 import io
11 import itertools
12 import logging
13 from multiprocessing import Manager, freeze_support
14 import os
15 from pathlib import Path
16 import pickle
17 import regex as re
18 import signal
19 import sys
20 import tempfile
21 import tokenize
22 import traceback
23 from typing import (
24     Any,
25     Callable,
26     Collection,
27     Dict,
28     Generator,
29     Generic,
30     Iterable,
31     Iterator,
32     List,
33     Optional,
34     Pattern,
35     Sequence,
36     Set,
37     Sized,
38     Tuple,
39     Type,
40     TypeVar,
41     Union,
42     cast,
43     TYPE_CHECKING,
44 )
45 from typing_extensions import Final
46 from mypy_extensions import mypyc_attr
47
48 from appdirs import user_cache_dir
49 from dataclasses import dataclass, field, replace
50 import click
51 import toml
52 from typed_ast import ast3, ast27
53 from pathspec import PathSpec
54
55 # lib2to3 fork
56 from blib2to3.pytree import Node, Leaf, type_repr
57 from blib2to3 import pygram, pytree
58 from blib2to3.pgen2 import driver, token
59 from blib2to3.pgen2.grammar import Grammar
60 from blib2to3.pgen2.parse import ParseError
61
62 from _black_version import version as __version__
63
64 if TYPE_CHECKING:
65     import colorama  # noqa: F401
66
67 DEFAULT_LINE_LENGTH = 88
68 DEFAULT_EXCLUDES = r"/(\.direnv|\.eggs|\.git|\.hg|\.mypy_cache|\.nox|\.tox|\.venv|\.svn|_build|buck-out|build|dist)/"  # noqa: B950
69 DEFAULT_INCLUDES = r"\.pyi?$"
70 CACHE_DIR = Path(user_cache_dir("black", version=__version__))
71
72 STRING_PREFIX_CHARS: Final = "furbFURB"  # All possible string prefix characters.
73
74
75 # types
76 FileContent = str
77 Encoding = str
78 NewLine = str
79 Depth = int
80 NodeType = int
81 ParserState = int
82 LeafID = int
83 StringID = int
84 Priority = int
85 Index = int
86 LN = Union[Leaf, Node]
87 Transformer = Callable[["Line", Collection["Feature"]], Iterator["Line"]]
88 Timestamp = float
89 FileSize = int
90 CacheInfo = Tuple[Timestamp, FileSize]
91 Cache = Dict[Path, CacheInfo]
92 out = partial(click.secho, bold=True, err=True)
93 err = partial(click.secho, fg="red", err=True)
94
95 pygram.initialize(CACHE_DIR)
96 syms = pygram.python_symbols
97
98
99 class NothingChanged(UserWarning):
100     """Raised when reformatted code is the same as source."""
101
102
103 class CannotTransform(Exception):
104     """Base class for errors raised by Transformers."""
105
106
107 class CannotSplit(CannotTransform):
108     """A readable split that fits the allotted line length is impossible."""
109
110
111 class InvalidInput(ValueError):
112     """Raised when input source code fails all parse attempts."""
113
114
115 class BracketMatchError(KeyError):
116     """Raised when an opening bracket is unable to be matched to a closing bracket."""
117
118
119 T = TypeVar("T")
120 E = TypeVar("E", bound=Exception)
121
122
123 class Ok(Generic[T]):
124     def __init__(self, value: T) -> None:
125         self._value = value
126
127     def ok(self) -> T:
128         return self._value
129
130
131 class Err(Generic[E]):
132     def __init__(self, e: E) -> None:
133         self._e = e
134
135     def err(self) -> E:
136         return self._e
137
138
139 # The 'Result' return type is used to implement an error-handling model heavily
140 # influenced by that used by the Rust programming language
141 # (see https://doc.rust-lang.org/book/ch09-00-error-handling.html).
142 Result = Union[Ok[T], Err[E]]
143 TResult = Result[T, CannotTransform]  # (T)ransform Result
144 TMatchResult = TResult[Index]
145
146
147 class WriteBack(Enum):
148     NO = 0
149     YES = 1
150     DIFF = 2
151     CHECK = 3
152     COLOR_DIFF = 4
153
154     @classmethod
155     def from_configuration(
156         cls, *, check: bool, diff: bool, color: bool = False
157     ) -> "WriteBack":
158         if check and not diff:
159             return cls.CHECK
160
161         if diff and color:
162             return cls.COLOR_DIFF
163
164         return cls.DIFF if diff else cls.YES
165
166
167 class Changed(Enum):
168     NO = 0
169     CACHED = 1
170     YES = 2
171
172
173 class TargetVersion(Enum):
174     PY27 = 2
175     PY33 = 3
176     PY34 = 4
177     PY35 = 5
178     PY36 = 6
179     PY37 = 7
180     PY38 = 8
181
182     def is_python2(self) -> bool:
183         return self is TargetVersion.PY27
184
185
186 PY36_VERSIONS = {TargetVersion.PY36, TargetVersion.PY37, TargetVersion.PY38}
187
188
189 class Feature(Enum):
190     # All string literals are unicode
191     UNICODE_LITERALS = 1
192     F_STRINGS = 2
193     NUMERIC_UNDERSCORES = 3
194     TRAILING_COMMA_IN_CALL = 4
195     TRAILING_COMMA_IN_DEF = 5
196     # The following two feature-flags are mutually exclusive, and exactly one should be
197     # set for every version of python.
198     ASYNC_IDENTIFIERS = 6
199     ASYNC_KEYWORDS = 7
200     ASSIGNMENT_EXPRESSIONS = 8
201     POS_ONLY_ARGUMENTS = 9
202     FORCE_OPTIONAL_PARENTHESES = 50
203
204
205 VERSION_TO_FEATURES: Dict[TargetVersion, Set[Feature]] = {
206     TargetVersion.PY27: {Feature.ASYNC_IDENTIFIERS},
207     TargetVersion.PY33: {Feature.UNICODE_LITERALS, Feature.ASYNC_IDENTIFIERS},
208     TargetVersion.PY34: {Feature.UNICODE_LITERALS, Feature.ASYNC_IDENTIFIERS},
209     TargetVersion.PY35: {
210         Feature.UNICODE_LITERALS,
211         Feature.TRAILING_COMMA_IN_CALL,
212         Feature.ASYNC_IDENTIFIERS,
213     },
214     TargetVersion.PY36: {
215         Feature.UNICODE_LITERALS,
216         Feature.F_STRINGS,
217         Feature.NUMERIC_UNDERSCORES,
218         Feature.TRAILING_COMMA_IN_CALL,
219         Feature.TRAILING_COMMA_IN_DEF,
220         Feature.ASYNC_IDENTIFIERS,
221     },
222     TargetVersion.PY37: {
223         Feature.UNICODE_LITERALS,
224         Feature.F_STRINGS,
225         Feature.NUMERIC_UNDERSCORES,
226         Feature.TRAILING_COMMA_IN_CALL,
227         Feature.TRAILING_COMMA_IN_DEF,
228         Feature.ASYNC_KEYWORDS,
229     },
230     TargetVersion.PY38: {
231         Feature.UNICODE_LITERALS,
232         Feature.F_STRINGS,
233         Feature.NUMERIC_UNDERSCORES,
234         Feature.TRAILING_COMMA_IN_CALL,
235         Feature.TRAILING_COMMA_IN_DEF,
236         Feature.ASYNC_KEYWORDS,
237         Feature.ASSIGNMENT_EXPRESSIONS,
238         Feature.POS_ONLY_ARGUMENTS,
239     },
240 }
241
242
243 @dataclass
244 class Mode:
245     target_versions: Set[TargetVersion] = field(default_factory=set)
246     line_length: int = DEFAULT_LINE_LENGTH
247     string_normalization: bool = True
248     experimental_string_processing: bool = False
249     is_pyi: bool = False
250
251     def get_cache_key(self) -> str:
252         if self.target_versions:
253             version_str = ",".join(
254                 str(version.value)
255                 for version in sorted(self.target_versions, key=lambda v: v.value)
256             )
257         else:
258             version_str = "-"
259         parts = [
260             version_str,
261             str(self.line_length),
262             str(int(self.string_normalization)),
263             str(int(self.is_pyi)),
264         ]
265         return ".".join(parts)
266
267
268 # Legacy name, left for integrations.
269 FileMode = Mode
270
271
272 def supports_feature(target_versions: Set[TargetVersion], feature: Feature) -> bool:
273     return all(feature in VERSION_TO_FEATURES[version] for version in target_versions)
274
275
276 def find_pyproject_toml(path_search_start: Iterable[str]) -> Optional[str]:
277     """Find the absolute filepath to a pyproject.toml if it exists"""
278     path_project_root = find_project_root(path_search_start)
279     path_pyproject_toml = path_project_root / "pyproject.toml"
280     return str(path_pyproject_toml) if path_pyproject_toml.is_file() else None
281
282
283 def parse_pyproject_toml(path_config: str) -> Dict[str, Any]:
284     """Parse a pyproject toml file, pulling out relevant parts for Black
285
286     If parsing fails, will raise a toml.TomlDecodeError
287     """
288     pyproject_toml = toml.load(path_config)
289     config = pyproject_toml.get("tool", {}).get("black", {})
290     return {k.replace("--", "").replace("-", "_"): v for k, v in config.items()}
291
292
293 def read_pyproject_toml(
294     ctx: click.Context, param: click.Parameter, value: Optional[str]
295 ) -> Optional[str]:
296     """Inject Black configuration from "pyproject.toml" into defaults in `ctx`.
297
298     Returns the path to a successfully found and read configuration file, None
299     otherwise.
300     """
301     if not value:
302         value = find_pyproject_toml(ctx.params.get("src", ()))
303         if value is None:
304             return None
305
306     try:
307         config = parse_pyproject_toml(value)
308     except (toml.TomlDecodeError, OSError) as e:
309         raise click.FileError(
310             filename=value, hint=f"Error reading configuration file: {e}"
311         )
312
313     if not config:
314         return None
315     else:
316         # Sanitize the values to be Click friendly. For more information please see:
317         # https://github.com/psf/black/issues/1458
318         # https://github.com/pallets/click/issues/1567
319         config = {
320             k: str(v) if not isinstance(v, (list, dict)) else v
321             for k, v in config.items()
322         }
323
324     target_version = config.get("target_version")
325     if target_version is not None and not isinstance(target_version, list):
326         raise click.BadOptionUsage(
327             "target-version", "Config key target-version must be a list"
328         )
329
330     default_map: Dict[str, Any] = {}
331     if ctx.default_map:
332         default_map.update(ctx.default_map)
333     default_map.update(config)
334
335     ctx.default_map = default_map
336     return value
337
338
339 def target_version_option_callback(
340     c: click.Context, p: Union[click.Option, click.Parameter], v: Tuple[str, ...]
341 ) -> List[TargetVersion]:
342     """Compute the target versions from a --target-version flag.
343
344     This is its own function because mypy couldn't infer the type correctly
345     when it was a lambda, causing mypyc trouble.
346     """
347     return [TargetVersion[val.upper()] for val in v]
348
349
350 @click.command(context_settings=dict(help_option_names=["-h", "--help"]))
351 @click.option("-c", "--code", type=str, help="Format the code passed in as a string.")
352 @click.option(
353     "-l",
354     "--line-length",
355     type=int,
356     default=DEFAULT_LINE_LENGTH,
357     help="How many characters per line to allow.",
358     show_default=True,
359 )
360 @click.option(
361     "-t",
362     "--target-version",
363     type=click.Choice([v.name.lower() for v in TargetVersion]),
364     callback=target_version_option_callback,
365     multiple=True,
366     help=(
367         "Python versions that should be supported by Black's output. [default: per-file"
368         " auto-detection]"
369     ),
370 )
371 @click.option(
372     "--pyi",
373     is_flag=True,
374     help=(
375         "Format all input files like typing stubs regardless of file extension (useful"
376         " when piping source on standard input)."
377     ),
378 )
379 @click.option(
380     "-S",
381     "--skip-string-normalization",
382     is_flag=True,
383     help="Don't normalize string quotes or prefixes.",
384 )
385 @click.option(
386     "--experimental-string-processing",
387     is_flag=True,
388     hidden=True,
389     help=(
390         "Experimental option that performs more normalization on string literals."
391         " Currently disabled because it leads to some crashes."
392     ),
393 )
394 @click.option(
395     "--check",
396     is_flag=True,
397     help=(
398         "Don't write the files back, just return the status.  Return code 0 means"
399         " nothing would change.  Return code 1 means some files would be reformatted."
400         " Return code 123 means there was an internal error."
401     ),
402 )
403 @click.option(
404     "--diff",
405     is_flag=True,
406     help="Don't write the files back, just output a diff for each file on stdout.",
407 )
408 @click.option(
409     "--color/--no-color",
410     is_flag=True,
411     help="Show colored diff. Only applies when `--diff` is given.",
412 )
413 @click.option(
414     "--fast/--safe",
415     is_flag=True,
416     help="If --fast given, skip temporary sanity checks. [default: --safe]",
417 )
418 @click.option(
419     "--include",
420     type=str,
421     default=DEFAULT_INCLUDES,
422     help=(
423         "A regular expression that matches files and directories that should be"
424         " included on recursive searches.  An empty value means all files are included"
425         " regardless of the name.  Use forward slashes for directories on all platforms"
426         " (Windows, too).  Exclusions are calculated first, inclusions later."
427     ),
428     show_default=True,
429 )
430 @click.option(
431     "--exclude",
432     type=str,
433     default=DEFAULT_EXCLUDES,
434     help=(
435         "A regular expression that matches files and directories that should be"
436         " excluded on recursive searches.  An empty value means no paths are excluded."
437         " Use forward slashes for directories on all platforms (Windows, too). "
438         " Exclusions are calculated first, inclusions later."
439     ),
440     show_default=True,
441 )
442 @click.option(
443     "--force-exclude",
444     type=str,
445     help=(
446         "Like --exclude, but files and directories matching this regex will be "
447         "excluded even when they are passed explicitly as arguments"
448     ),
449 )
450 @click.option(
451     "-q",
452     "--quiet",
453     is_flag=True,
454     help=(
455         "Don't emit non-error messages to stderr. Errors are still emitted; silence"
456         " those with 2>/dev/null."
457     ),
458 )
459 @click.option(
460     "-v",
461     "--verbose",
462     is_flag=True,
463     help=(
464         "Also emit messages to stderr about files that were not changed or were ignored"
465         " due to --exclude=."
466     ),
467 )
468 @click.version_option(version=__version__)
469 @click.argument(
470     "src",
471     nargs=-1,
472     type=click.Path(
473         exists=True, file_okay=True, dir_okay=True, readable=True, allow_dash=True
474     ),
475     is_eager=True,
476 )
477 @click.option(
478     "--config",
479     type=click.Path(
480         exists=True,
481         file_okay=True,
482         dir_okay=False,
483         readable=True,
484         allow_dash=False,
485         path_type=str,
486     ),
487     is_eager=True,
488     callback=read_pyproject_toml,
489     help="Read configuration from FILE path.",
490 )
491 @click.pass_context
492 def main(
493     ctx: click.Context,
494     code: Optional[str],
495     line_length: int,
496     target_version: List[TargetVersion],
497     check: bool,
498     diff: bool,
499     color: bool,
500     fast: bool,
501     pyi: bool,
502     skip_string_normalization: bool,
503     experimental_string_processing: bool,
504     quiet: bool,
505     verbose: bool,
506     include: str,
507     exclude: str,
508     force_exclude: Optional[str],
509     src: Tuple[str, ...],
510     config: Optional[str],
511 ) -> None:
512     """The uncompromising code formatter."""
513     write_back = WriteBack.from_configuration(check=check, diff=diff, color=color)
514     if target_version:
515         versions = set(target_version)
516     else:
517         # We'll autodetect later.
518         versions = set()
519     mode = Mode(
520         target_versions=versions,
521         line_length=line_length,
522         is_pyi=pyi,
523         string_normalization=not skip_string_normalization,
524         experimental_string_processing=experimental_string_processing,
525     )
526     if config and verbose:
527         out(f"Using configuration from {config}.", bold=False, fg="blue")
528     if code is not None:
529         print(format_str(code, mode=mode))
530         ctx.exit(0)
531     report = Report(check=check, diff=diff, quiet=quiet, verbose=verbose)
532     sources = get_sources(
533         ctx=ctx,
534         src=src,
535         quiet=quiet,
536         verbose=verbose,
537         include=include,
538         exclude=exclude,
539         force_exclude=force_exclude,
540         report=report,
541     )
542
543     path_empty(
544         sources,
545         "No Python files are present to be formatted. Nothing to do 😴",
546         quiet,
547         verbose,
548         ctx,
549     )
550
551     if len(sources) == 1:
552         reformat_one(
553             src=sources.pop(),
554             fast=fast,
555             write_back=write_back,
556             mode=mode,
557             report=report,
558         )
559     else:
560         reformat_many(
561             sources=sources, fast=fast, write_back=write_back, mode=mode, report=report
562         )
563
564     if verbose or not quiet:
565         out("Oh no! 💥 💔 💥" if report.return_code else "All done! ✨ 🍰 ✨")
566         click.secho(str(report), err=True)
567     ctx.exit(report.return_code)
568
569
570 def get_sources(
571     *,
572     ctx: click.Context,
573     src: Tuple[str, ...],
574     quiet: bool,
575     verbose: bool,
576     include: str,
577     exclude: str,
578     force_exclude: Optional[str],
579     report: "Report",
580 ) -> Set[Path]:
581     """Compute the set of files to be formatted."""
582     try:
583         include_regex = re_compile_maybe_verbose(include)
584     except re.error:
585         err(f"Invalid regular expression for include given: {include!r}")
586         ctx.exit(2)
587     try:
588         exclude_regex = re_compile_maybe_verbose(exclude)
589     except re.error:
590         err(f"Invalid regular expression for exclude given: {exclude!r}")
591         ctx.exit(2)
592     try:
593         force_exclude_regex = (
594             re_compile_maybe_verbose(force_exclude) if force_exclude else None
595         )
596     except re.error:
597         err(f"Invalid regular expression for force_exclude given: {force_exclude!r}")
598         ctx.exit(2)
599
600     root = find_project_root(src)
601     sources: Set[Path] = set()
602     path_empty(src, "No Path provided. Nothing to do 😴", quiet, verbose, ctx)
603     gitignore = get_gitignore(root)
604
605     for s in src:
606         p = Path(s)
607         if p.is_dir():
608             sources.update(
609                 gen_python_files(
610                     p.iterdir(),
611                     root,
612                     include_regex,
613                     exclude_regex,
614                     force_exclude_regex,
615                     report,
616                     gitignore,
617                 )
618             )
619         elif s == "-":
620             sources.add(p)
621         elif p.is_file():
622             normalized_path = normalize_path_maybe_ignore(p, root, report)
623             if normalized_path is None:
624                 continue
625
626             normalized_path = "/" + normalized_path
627             # Hard-exclude any files that matches the `--force-exclude` regex.
628             if force_exclude_regex:
629                 force_exclude_match = force_exclude_regex.search(normalized_path)
630             else:
631                 force_exclude_match = None
632             if force_exclude_match and force_exclude_match.group(0):
633                 report.path_ignored(p, "matches the --force-exclude regular expression")
634                 continue
635
636             sources.add(p)
637         else:
638             err(f"invalid path: {s}")
639     return sources
640
641
642 def path_empty(
643     src: Sized, msg: str, quiet: bool, verbose: bool, ctx: click.Context
644 ) -> None:
645     """
646     Exit if there is no `src` provided for formatting
647     """
648     if not src and (verbose or not quiet):
649         out(msg)
650         ctx.exit(0)
651
652
653 def reformat_one(
654     src: Path, fast: bool, write_back: WriteBack, mode: Mode, report: "Report"
655 ) -> None:
656     """Reformat a single file under `src` without spawning child processes.
657
658     `fast`, `write_back`, and `mode` options are passed to
659     :func:`format_file_in_place` or :func:`format_stdin_to_stdout`.
660     """
661     try:
662         changed = Changed.NO
663         if not src.is_file() and str(src) == "-":
664             if format_stdin_to_stdout(fast=fast, write_back=write_back, mode=mode):
665                 changed = Changed.YES
666         else:
667             cache: Cache = {}
668             if write_back not in (WriteBack.DIFF, WriteBack.COLOR_DIFF):
669                 cache = read_cache(mode)
670                 res_src = src.resolve()
671                 if res_src in cache and cache[res_src] == get_cache_info(res_src):
672                     changed = Changed.CACHED
673             if changed is not Changed.CACHED and format_file_in_place(
674                 src, fast=fast, write_back=write_back, mode=mode
675             ):
676                 changed = Changed.YES
677             if (write_back is WriteBack.YES and changed is not Changed.CACHED) or (
678                 write_back is WriteBack.CHECK and changed is Changed.NO
679             ):
680                 write_cache(cache, [src], mode)
681         report.done(src, changed)
682     except Exception as exc:
683         if report.verbose:
684             traceback.print_exc()
685         report.failed(src, str(exc))
686
687
688 def reformat_many(
689     sources: Set[Path], fast: bool, write_back: WriteBack, mode: Mode, report: "Report"
690 ) -> None:
691     """Reformat multiple files using a ProcessPoolExecutor."""
692     executor: Executor
693     loop = asyncio.get_event_loop()
694     worker_count = os.cpu_count()
695     if sys.platform == "win32":
696         # Work around https://bugs.python.org/issue26903
697         worker_count = min(worker_count, 61)
698     try:
699         executor = ProcessPoolExecutor(max_workers=worker_count)
700     except (ImportError, OSError):
701         # we arrive here if the underlying system does not support multi-processing
702         # like in AWS Lambda or Termux, in which case we gracefully fallback to
703         # a ThreadPollExecutor with just a single worker (more workers would not do us
704         # any good due to the Global Interpreter Lock)
705         executor = ThreadPoolExecutor(max_workers=1)
706
707     try:
708         loop.run_until_complete(
709             schedule_formatting(
710                 sources=sources,
711                 fast=fast,
712                 write_back=write_back,
713                 mode=mode,
714                 report=report,
715                 loop=loop,
716                 executor=executor,
717             )
718         )
719     finally:
720         shutdown(loop)
721         if executor is not None:
722             executor.shutdown()
723
724
725 async def schedule_formatting(
726     sources: Set[Path],
727     fast: bool,
728     write_back: WriteBack,
729     mode: Mode,
730     report: "Report",
731     loop: asyncio.AbstractEventLoop,
732     executor: Executor,
733 ) -> None:
734     """Run formatting of `sources` in parallel using the provided `executor`.
735
736     (Use ProcessPoolExecutors for actual parallelism.)
737
738     `write_back`, `fast`, and `mode` options are passed to
739     :func:`format_file_in_place`.
740     """
741     cache: Cache = {}
742     if write_back not in (WriteBack.DIFF, WriteBack.COLOR_DIFF):
743         cache = read_cache(mode)
744         sources, cached = filter_cached(cache, sources)
745         for src in sorted(cached):
746             report.done(src, Changed.CACHED)
747     if not sources:
748         return
749
750     cancelled = []
751     sources_to_cache = []
752     lock = None
753     if write_back in (WriteBack.DIFF, WriteBack.COLOR_DIFF):
754         # For diff output, we need locks to ensure we don't interleave output
755         # from different processes.
756         manager = Manager()
757         lock = manager.Lock()
758     tasks = {
759         asyncio.ensure_future(
760             loop.run_in_executor(
761                 executor, format_file_in_place, src, fast, mode, write_back, lock
762             )
763         ): src
764         for src in sorted(sources)
765     }
766     pending: Iterable["asyncio.Future[bool]"] = tasks.keys()
767     try:
768         loop.add_signal_handler(signal.SIGINT, cancel, pending)
769         loop.add_signal_handler(signal.SIGTERM, cancel, pending)
770     except NotImplementedError:
771         # There are no good alternatives for these on Windows.
772         pass
773     while pending:
774         done, _ = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
775         for task in done:
776             src = tasks.pop(task)
777             if task.cancelled():
778                 cancelled.append(task)
779             elif task.exception():
780                 report.failed(src, str(task.exception()))
781             else:
782                 changed = Changed.YES if task.result() else Changed.NO
783                 # If the file was written back or was successfully checked as
784                 # well-formatted, store this information in the cache.
785                 if write_back is WriteBack.YES or (
786                     write_back is WriteBack.CHECK and changed is Changed.NO
787                 ):
788                     sources_to_cache.append(src)
789                 report.done(src, changed)
790     if cancelled:
791         await asyncio.gather(*cancelled, loop=loop, return_exceptions=True)
792     if sources_to_cache:
793         write_cache(cache, sources_to_cache, mode)
794
795
796 def format_file_in_place(
797     src: Path,
798     fast: bool,
799     mode: Mode,
800     write_back: WriteBack = WriteBack.NO,
801     lock: Any = None,  # multiprocessing.Manager().Lock() is some crazy proxy
802 ) -> bool:
803     """Format file under `src` path. Return True if changed.
804
805     If `write_back` is DIFF, write a diff to stdout. If it is YES, write reformatted
806     code to the file.
807     `mode` and `fast` options are passed to :func:`format_file_contents`.
808     """
809     if src.suffix == ".pyi":
810         mode = replace(mode, is_pyi=True)
811
812     then = datetime.utcfromtimestamp(src.stat().st_mtime)
813     with open(src, "rb") as buf:
814         src_contents, encoding, newline = decode_bytes(buf.read())
815     try:
816         dst_contents = format_file_contents(src_contents, fast=fast, mode=mode)
817     except NothingChanged:
818         return False
819
820     if write_back == WriteBack.YES:
821         with open(src, "w", encoding=encoding, newline=newline) as f:
822             f.write(dst_contents)
823     elif write_back in (WriteBack.DIFF, WriteBack.COLOR_DIFF):
824         now = datetime.utcnow()
825         src_name = f"{src}\t{then} +0000"
826         dst_name = f"{src}\t{now} +0000"
827         diff_contents = diff(src_contents, dst_contents, src_name, dst_name)
828
829         if write_back == write_back.COLOR_DIFF:
830             diff_contents = color_diff(diff_contents)
831
832         with lock or nullcontext():
833             f = io.TextIOWrapper(
834                 sys.stdout.buffer,
835                 encoding=encoding,
836                 newline=newline,
837                 write_through=True,
838             )
839             f = wrap_stream_for_windows(f)
840             f.write(diff_contents)
841             f.detach()
842
843     return True
844
845
846 def color_diff(contents: str) -> str:
847     """Inject the ANSI color codes to the diff."""
848     lines = contents.split("\n")
849     for i, line in enumerate(lines):
850         if line.startswith("+++") or line.startswith("---"):
851             line = "\033[1;37m" + line + "\033[0m"  # bold white, reset
852         if line.startswith("@@"):
853             line = "\033[36m" + line + "\033[0m"  # cyan, reset
854         if line.startswith("+"):
855             line = "\033[32m" + line + "\033[0m"  # green, reset
856         elif line.startswith("-"):
857             line = "\033[31m" + line + "\033[0m"  # red, reset
858         lines[i] = line
859     return "\n".join(lines)
860
861
862 def wrap_stream_for_windows(
863     f: io.TextIOWrapper,
864 ) -> Union[io.TextIOWrapper, "colorama.AnsiToWin32.AnsiToWin32"]:
865     """
866     Wrap the stream in colorama's wrap_stream so colors are shown on Windows.
867
868     If `colorama` is not found, then no change is made. If `colorama` does
869     exist, then it handles the logic to determine whether or not to change
870     things.
871     """
872     try:
873         from colorama import initialise
874
875         # We set `strip=False` so that we can don't have to modify
876         # test_express_diff_with_color.
877         f = initialise.wrap_stream(
878             f, convert=None, strip=False, autoreset=False, wrap=True
879         )
880
881         # wrap_stream returns a `colorama.AnsiToWin32.AnsiToWin32` object
882         # which does not have a `detach()` method. So we fake one.
883         f.detach = lambda *args, **kwargs: None  # type: ignore
884     except ImportError:
885         pass
886
887     return f
888
889
890 def format_stdin_to_stdout(
891     fast: bool, *, write_back: WriteBack = WriteBack.NO, mode: Mode
892 ) -> bool:
893     """Format file on stdin. Return True if changed.
894
895     If `write_back` is YES, write reformatted code back to stdout. If it is DIFF,
896     write a diff to stdout. The `mode` argument is passed to
897     :func:`format_file_contents`.
898     """
899     then = datetime.utcnow()
900     src, encoding, newline = decode_bytes(sys.stdin.buffer.read())
901     dst = src
902     try:
903         dst = format_file_contents(src, fast=fast, mode=mode)
904         return True
905
906     except NothingChanged:
907         return False
908
909     finally:
910         f = io.TextIOWrapper(
911             sys.stdout.buffer, encoding=encoding, newline=newline, write_through=True
912         )
913         if write_back == WriteBack.YES:
914             f.write(dst)
915         elif write_back in (WriteBack.DIFF, WriteBack.COLOR_DIFF):
916             now = datetime.utcnow()
917             src_name = f"STDIN\t{then} +0000"
918             dst_name = f"STDOUT\t{now} +0000"
919             d = diff(src, dst, src_name, dst_name)
920             if write_back == WriteBack.COLOR_DIFF:
921                 d = color_diff(d)
922                 f = wrap_stream_for_windows(f)
923             f.write(d)
924         f.detach()
925
926
927 def format_file_contents(src_contents: str, *, fast: bool, mode: Mode) -> FileContent:
928     """Reformat contents a file and return new contents.
929
930     If `fast` is False, additionally confirm that the reformatted code is
931     valid by calling :func:`assert_equivalent` and :func:`assert_stable` on it.
932     `mode` is passed to :func:`format_str`.
933     """
934     if not src_contents.strip():
935         raise NothingChanged
936
937     dst_contents = format_str(src_contents, mode=mode)
938     if src_contents == dst_contents:
939         raise NothingChanged
940
941     if not fast:
942         assert_equivalent(src_contents, dst_contents)
943         assert_stable(src_contents, dst_contents, mode=mode)
944     return dst_contents
945
946
947 def format_str(src_contents: str, *, mode: Mode) -> FileContent:
948     """Reformat a string and return new contents.
949
950     `mode` determines formatting options, such as how many characters per line are
951     allowed.  Example:
952
953     >>> import black
954     >>> print(black.format_str("def f(arg:str='')->None:...", mode=Mode()))
955     def f(arg: str = "") -> None:
956         ...
957
958     A more complex example:
959
960     >>> print(
961     ...   black.format_str(
962     ...     "def f(arg:str='')->None: hey",
963     ...     mode=black.Mode(
964     ...       target_versions={black.TargetVersion.PY36},
965     ...       line_length=10,
966     ...       string_normalization=False,
967     ...       is_pyi=False,
968     ...     ),
969     ...   ),
970     ... )
971     def f(
972         arg: str = '',
973     ) -> None:
974         hey
975
976     """
977     src_node = lib2to3_parse(src_contents.lstrip(), mode.target_versions)
978     dst_contents = []
979     future_imports = get_future_imports(src_node)
980     if mode.target_versions:
981         versions = mode.target_versions
982     else:
983         versions = detect_target_versions(src_node)
984     normalize_fmt_off(src_node)
985     lines = LineGenerator(
986         remove_u_prefix="unicode_literals" in future_imports
987         or supports_feature(versions, Feature.UNICODE_LITERALS),
988         is_pyi=mode.is_pyi,
989         normalize_strings=mode.string_normalization,
990     )
991     elt = EmptyLineTracker(is_pyi=mode.is_pyi)
992     empty_line = Line()
993     after = 0
994     split_line_features = {
995         feature
996         for feature in {Feature.TRAILING_COMMA_IN_CALL, Feature.TRAILING_COMMA_IN_DEF}
997         if supports_feature(versions, feature)
998     }
999     for current_line in lines.visit(src_node):
1000         dst_contents.append(str(empty_line) * after)
1001         before, after = elt.maybe_empty_lines(current_line)
1002         dst_contents.append(str(empty_line) * before)
1003         for line in transform_line(
1004             current_line, mode=mode, features=split_line_features
1005         ):
1006             dst_contents.append(str(line))
1007     return "".join(dst_contents)
1008
1009
1010 def decode_bytes(src: bytes) -> Tuple[FileContent, Encoding, NewLine]:
1011     """Return a tuple of (decoded_contents, encoding, newline).
1012
1013     `newline` is either CRLF or LF but `decoded_contents` is decoded with
1014     universal newlines (i.e. only contains LF).
1015     """
1016     srcbuf = io.BytesIO(src)
1017     encoding, lines = tokenize.detect_encoding(srcbuf.readline)
1018     if not lines:
1019         return "", encoding, "\n"
1020
1021     newline = "\r\n" if b"\r\n" == lines[0][-2:] else "\n"
1022     srcbuf.seek(0)
1023     with io.TextIOWrapper(srcbuf, encoding) as tiow:
1024         return tiow.read(), encoding, newline
1025
1026
1027 def get_grammars(target_versions: Set[TargetVersion]) -> List[Grammar]:
1028     if not target_versions:
1029         # No target_version specified, so try all grammars.
1030         return [
1031             # Python 3.7+
1032             pygram.python_grammar_no_print_statement_no_exec_statement_async_keywords,
1033             # Python 3.0-3.6
1034             pygram.python_grammar_no_print_statement_no_exec_statement,
1035             # Python 2.7 with future print_function import
1036             pygram.python_grammar_no_print_statement,
1037             # Python 2.7
1038             pygram.python_grammar,
1039         ]
1040
1041     if all(version.is_python2() for version in target_versions):
1042         # Python 2-only code, so try Python 2 grammars.
1043         return [
1044             # Python 2.7 with future print_function import
1045             pygram.python_grammar_no_print_statement,
1046             # Python 2.7
1047             pygram.python_grammar,
1048         ]
1049
1050     # Python 3-compatible code, so only try Python 3 grammar.
1051     grammars = []
1052     # If we have to parse both, try to parse async as a keyword first
1053     if not supports_feature(target_versions, Feature.ASYNC_IDENTIFIERS):
1054         # Python 3.7+
1055         grammars.append(
1056             pygram.python_grammar_no_print_statement_no_exec_statement_async_keywords
1057         )
1058     if not supports_feature(target_versions, Feature.ASYNC_KEYWORDS):
1059         # Python 3.0-3.6
1060         grammars.append(pygram.python_grammar_no_print_statement_no_exec_statement)
1061     # At least one of the above branches must have been taken, because every Python
1062     # version has exactly one of the two 'ASYNC_*' flags
1063     return grammars
1064
1065
1066 def lib2to3_parse(src_txt: str, target_versions: Iterable[TargetVersion] = ()) -> Node:
1067     """Given a string with source, return the lib2to3 Node."""
1068     if not src_txt.endswith("\n"):
1069         src_txt += "\n"
1070
1071     for grammar in get_grammars(set(target_versions)):
1072         drv = driver.Driver(grammar, pytree.convert)
1073         try:
1074             result = drv.parse_string(src_txt, True)
1075             break
1076
1077         except ParseError as pe:
1078             lineno, column = pe.context[1]
1079             lines = src_txt.splitlines()
1080             try:
1081                 faulty_line = lines[lineno - 1]
1082             except IndexError:
1083                 faulty_line = "<line number missing in source>"
1084             exc = InvalidInput(f"Cannot parse: {lineno}:{column}: {faulty_line}")
1085     else:
1086         raise exc from None
1087
1088     if isinstance(result, Leaf):
1089         result = Node(syms.file_input, [result])
1090     return result
1091
1092
1093 def lib2to3_unparse(node: Node) -> str:
1094     """Given a lib2to3 node, return its string representation."""
1095     code = str(node)
1096     return code
1097
1098
1099 class Visitor(Generic[T]):
1100     """Basic lib2to3 visitor that yields things of type `T` on `visit()`."""
1101
1102     def visit(self, node: LN) -> Iterator[T]:
1103         """Main method to visit `node` and its children.
1104
1105         It tries to find a `visit_*()` method for the given `node.type`, like
1106         `visit_simple_stmt` for Node objects or `visit_INDENT` for Leaf objects.
1107         If no dedicated `visit_*()` method is found, chooses `visit_default()`
1108         instead.
1109
1110         Then yields objects of type `T` from the selected visitor.
1111         """
1112         if node.type < 256:
1113             name = token.tok_name[node.type]
1114         else:
1115             name = str(type_repr(node.type))
1116         # We explicitly branch on whether a visitor exists (instead of
1117         # using self.visit_default as the default arg to getattr) in order
1118         # to save needing to create a bound method object and so mypyc can
1119         # generate a native call to visit_default.
1120         visitf = getattr(self, f"visit_{name}", None)
1121         if visitf:
1122             yield from visitf(node)
1123         else:
1124             yield from self.visit_default(node)
1125
1126     def visit_default(self, node: LN) -> Iterator[T]:
1127         """Default `visit_*()` implementation. Recurses to children of `node`."""
1128         if isinstance(node, Node):
1129             for child in node.children:
1130                 yield from self.visit(child)
1131
1132
1133 @dataclass
1134 class DebugVisitor(Visitor[T]):
1135     tree_depth: int = 0
1136
1137     def visit_default(self, node: LN) -> Iterator[T]:
1138         indent = " " * (2 * self.tree_depth)
1139         if isinstance(node, Node):
1140             _type = type_repr(node.type)
1141             out(f"{indent}{_type}", fg="yellow")
1142             self.tree_depth += 1
1143             for child in node.children:
1144                 yield from self.visit(child)
1145
1146             self.tree_depth -= 1
1147             out(f"{indent}/{_type}", fg="yellow", bold=False)
1148         else:
1149             _type = token.tok_name.get(node.type, str(node.type))
1150             out(f"{indent}{_type}", fg="blue", nl=False)
1151             if node.prefix:
1152                 # We don't have to handle prefixes for `Node` objects since
1153                 # that delegates to the first child anyway.
1154                 out(f" {node.prefix!r}", fg="green", bold=False, nl=False)
1155             out(f" {node.value!r}", fg="blue", bold=False)
1156
1157     @classmethod
1158     def show(cls, code: Union[str, Leaf, Node]) -> None:
1159         """Pretty-print the lib2to3 AST of a given string of `code`.
1160
1161         Convenience method for debugging.
1162         """
1163         v: DebugVisitor[None] = DebugVisitor()
1164         if isinstance(code, str):
1165             code = lib2to3_parse(code)
1166         list(v.visit(code))
1167
1168
1169 WHITESPACE: Final = {token.DEDENT, token.INDENT, token.NEWLINE}
1170 STATEMENT: Final = {
1171     syms.if_stmt,
1172     syms.while_stmt,
1173     syms.for_stmt,
1174     syms.try_stmt,
1175     syms.except_clause,
1176     syms.with_stmt,
1177     syms.funcdef,
1178     syms.classdef,
1179 }
1180 STANDALONE_COMMENT: Final = 153
1181 token.tok_name[STANDALONE_COMMENT] = "STANDALONE_COMMENT"
1182 LOGIC_OPERATORS: Final = {"and", "or"}
1183 COMPARATORS: Final = {
1184     token.LESS,
1185     token.GREATER,
1186     token.EQEQUAL,
1187     token.NOTEQUAL,
1188     token.LESSEQUAL,
1189     token.GREATEREQUAL,
1190 }
1191 MATH_OPERATORS: Final = {
1192     token.VBAR,
1193     token.CIRCUMFLEX,
1194     token.AMPER,
1195     token.LEFTSHIFT,
1196     token.RIGHTSHIFT,
1197     token.PLUS,
1198     token.MINUS,
1199     token.STAR,
1200     token.SLASH,
1201     token.DOUBLESLASH,
1202     token.PERCENT,
1203     token.AT,
1204     token.TILDE,
1205     token.DOUBLESTAR,
1206 }
1207 STARS: Final = {token.STAR, token.DOUBLESTAR}
1208 VARARGS_SPECIALS: Final = STARS | {token.SLASH}
1209 VARARGS_PARENTS: Final = {
1210     syms.arglist,
1211     syms.argument,  # double star in arglist
1212     syms.trailer,  # single argument to call
1213     syms.typedargslist,
1214     syms.varargslist,  # lambdas
1215 }
1216 UNPACKING_PARENTS: Final = {
1217     syms.atom,  # single element of a list or set literal
1218     syms.dictsetmaker,
1219     syms.listmaker,
1220     syms.testlist_gexp,
1221     syms.testlist_star_expr,
1222 }
1223 TEST_DESCENDANTS: Final = {
1224     syms.test,
1225     syms.lambdef,
1226     syms.or_test,
1227     syms.and_test,
1228     syms.not_test,
1229     syms.comparison,
1230     syms.star_expr,
1231     syms.expr,
1232     syms.xor_expr,
1233     syms.and_expr,
1234     syms.shift_expr,
1235     syms.arith_expr,
1236     syms.trailer,
1237     syms.term,
1238     syms.power,
1239 }
1240 ASSIGNMENTS: Final = {
1241     "=",
1242     "+=",
1243     "-=",
1244     "*=",
1245     "@=",
1246     "/=",
1247     "%=",
1248     "&=",
1249     "|=",
1250     "^=",
1251     "<<=",
1252     ">>=",
1253     "**=",
1254     "//=",
1255 }
1256 COMPREHENSION_PRIORITY: Final = 20
1257 COMMA_PRIORITY: Final = 18
1258 TERNARY_PRIORITY: Final = 16
1259 LOGIC_PRIORITY: Final = 14
1260 STRING_PRIORITY: Final = 12
1261 COMPARATOR_PRIORITY: Final = 10
1262 MATH_PRIORITIES: Final = {
1263     token.VBAR: 9,
1264     token.CIRCUMFLEX: 8,
1265     token.AMPER: 7,
1266     token.LEFTSHIFT: 6,
1267     token.RIGHTSHIFT: 6,
1268     token.PLUS: 5,
1269     token.MINUS: 5,
1270     token.STAR: 4,
1271     token.SLASH: 4,
1272     token.DOUBLESLASH: 4,
1273     token.PERCENT: 4,
1274     token.AT: 4,
1275     token.TILDE: 3,
1276     token.DOUBLESTAR: 2,
1277 }
1278 DOT_PRIORITY: Final = 1
1279
1280
1281 @dataclass
1282 class BracketTracker:
1283     """Keeps track of brackets on a line."""
1284
1285     depth: int = 0
1286     bracket_match: Dict[Tuple[Depth, NodeType], Leaf] = field(default_factory=dict)
1287     delimiters: Dict[LeafID, Priority] = field(default_factory=dict)
1288     previous: Optional[Leaf] = None
1289     _for_loop_depths: List[int] = field(default_factory=list)
1290     _lambda_argument_depths: List[int] = field(default_factory=list)
1291     invisible: List[Leaf] = field(default_factory=list)
1292
1293     def mark(self, leaf: Leaf) -> None:
1294         """Mark `leaf` with bracket-related metadata. Keep track of delimiters.
1295
1296         All leaves receive an int `bracket_depth` field that stores how deep
1297         within brackets a given leaf is. 0 means there are no enclosing brackets
1298         that started on this line.
1299
1300         If a leaf is itself a closing bracket, it receives an `opening_bracket`
1301         field that it forms a pair with. This is a one-directional link to
1302         avoid reference cycles.
1303
1304         If a leaf is a delimiter (a token on which Black can split the line if
1305         needed) and it's on depth 0, its `id()` is stored in the tracker's
1306         `delimiters` field.
1307         """
1308         if leaf.type == token.COMMENT:
1309             return
1310
1311         self.maybe_decrement_after_for_loop_variable(leaf)
1312         self.maybe_decrement_after_lambda_arguments(leaf)
1313         if leaf.type in CLOSING_BRACKETS:
1314             self.depth -= 1
1315             try:
1316                 opening_bracket = self.bracket_match.pop((self.depth, leaf.type))
1317             except KeyError as e:
1318                 raise BracketMatchError(
1319                     "Unable to match a closing bracket to the following opening"
1320                     f" bracket: {leaf}"
1321                 ) from e
1322             leaf.opening_bracket = opening_bracket
1323             if not leaf.value:
1324                 self.invisible.append(leaf)
1325         leaf.bracket_depth = self.depth
1326         if self.depth == 0:
1327             delim = is_split_before_delimiter(leaf, self.previous)
1328             if delim and self.previous is not None:
1329                 self.delimiters[id(self.previous)] = delim
1330             else:
1331                 delim = is_split_after_delimiter(leaf, self.previous)
1332                 if delim:
1333                     self.delimiters[id(leaf)] = delim
1334         if leaf.type in OPENING_BRACKETS:
1335             self.bracket_match[self.depth, BRACKET[leaf.type]] = leaf
1336             self.depth += 1
1337             if not leaf.value:
1338                 self.invisible.append(leaf)
1339         self.previous = leaf
1340         self.maybe_increment_lambda_arguments(leaf)
1341         self.maybe_increment_for_loop_variable(leaf)
1342
1343     def any_open_brackets(self) -> bool:
1344         """Return True if there is an yet unmatched open bracket on the line."""
1345         return bool(self.bracket_match)
1346
1347     def max_delimiter_priority(self, exclude: Iterable[LeafID] = ()) -> Priority:
1348         """Return the highest priority of a delimiter found on the line.
1349
1350         Values are consistent with what `is_split_*_delimiter()` return.
1351         Raises ValueError on no delimiters.
1352         """
1353         return max(v for k, v in self.delimiters.items() if k not in exclude)
1354
1355     def delimiter_count_with_priority(self, priority: Priority = 0) -> int:
1356         """Return the number of delimiters with the given `priority`.
1357
1358         If no `priority` is passed, defaults to max priority on the line.
1359         """
1360         if not self.delimiters:
1361             return 0
1362
1363         priority = priority or self.max_delimiter_priority()
1364         return sum(1 for p in self.delimiters.values() if p == priority)
1365
1366     def maybe_increment_for_loop_variable(self, leaf: Leaf) -> bool:
1367         """In a for loop, or comprehension, the variables are often unpacks.
1368
1369         To avoid splitting on the comma in this situation, increase the depth of
1370         tokens between `for` and `in`.
1371         """
1372         if leaf.type == token.NAME and leaf.value == "for":
1373             self.depth += 1
1374             self._for_loop_depths.append(self.depth)
1375             return True
1376
1377         return False
1378
1379     def maybe_decrement_after_for_loop_variable(self, leaf: Leaf) -> bool:
1380         """See `maybe_increment_for_loop_variable` above for explanation."""
1381         if (
1382             self._for_loop_depths
1383             and self._for_loop_depths[-1] == self.depth
1384             and leaf.type == token.NAME
1385             and leaf.value == "in"
1386         ):
1387             self.depth -= 1
1388             self._for_loop_depths.pop()
1389             return True
1390
1391         return False
1392
1393     def maybe_increment_lambda_arguments(self, leaf: Leaf) -> bool:
1394         """In a lambda expression, there might be more than one argument.
1395
1396         To avoid splitting on the comma in this situation, increase the depth of
1397         tokens between `lambda` and `:`.
1398         """
1399         if leaf.type == token.NAME and leaf.value == "lambda":
1400             self.depth += 1
1401             self._lambda_argument_depths.append(self.depth)
1402             return True
1403
1404         return False
1405
1406     def maybe_decrement_after_lambda_arguments(self, leaf: Leaf) -> bool:
1407         """See `maybe_increment_lambda_arguments` above for explanation."""
1408         if (
1409             self._lambda_argument_depths
1410             and self._lambda_argument_depths[-1] == self.depth
1411             and leaf.type == token.COLON
1412         ):
1413             self.depth -= 1
1414             self._lambda_argument_depths.pop()
1415             return True
1416
1417         return False
1418
1419     def get_open_lsqb(self) -> Optional[Leaf]:
1420         """Return the most recent opening square bracket (if any)."""
1421         return self.bracket_match.get((self.depth - 1, token.RSQB))
1422
1423
1424 @dataclass
1425 class Line:
1426     """Holds leaves and comments. Can be printed with `str(line)`."""
1427
1428     depth: int = 0
1429     leaves: List[Leaf] = field(default_factory=list)
1430     # keys ordered like `leaves`
1431     comments: Dict[LeafID, List[Leaf]] = field(default_factory=dict)
1432     bracket_tracker: BracketTracker = field(default_factory=BracketTracker)
1433     inside_brackets: bool = False
1434     should_explode: bool = False
1435
1436     def append(self, leaf: Leaf, preformatted: bool = False) -> None:
1437         """Add a new `leaf` to the end of the line.
1438
1439         Unless `preformatted` is True, the `leaf` will receive a new consistent
1440         whitespace prefix and metadata applied by :class:`BracketTracker`.
1441         Trailing commas are maybe removed, unpacked for loop variables are
1442         demoted from being delimiters.
1443
1444         Inline comments are put aside.
1445         """
1446         has_value = leaf.type in BRACKETS or bool(leaf.value.strip())
1447         if not has_value:
1448             return
1449
1450         if token.COLON == leaf.type and self.is_class_paren_empty:
1451             del self.leaves[-2:]
1452         if self.leaves and not preformatted:
1453             # Note: at this point leaf.prefix should be empty except for
1454             # imports, for which we only preserve newlines.
1455             leaf.prefix += whitespace(
1456                 leaf, complex_subscript=self.is_complex_subscript(leaf)
1457             )
1458         if self.inside_brackets or not preformatted:
1459             self.bracket_tracker.mark(leaf)
1460             if self.maybe_should_explode(leaf):
1461                 self.should_explode = True
1462         if not self.append_comment(leaf):
1463             self.leaves.append(leaf)
1464
1465     def append_safe(self, leaf: Leaf, preformatted: bool = False) -> None:
1466         """Like :func:`append()` but disallow invalid standalone comment structure.
1467
1468         Raises ValueError when any `leaf` is appended after a standalone comment
1469         or when a standalone comment is not the first leaf on the line.
1470         """
1471         if self.bracket_tracker.depth == 0:
1472             if self.is_comment:
1473                 raise ValueError("cannot append to standalone comments")
1474
1475             if self.leaves and leaf.type == STANDALONE_COMMENT:
1476                 raise ValueError(
1477                     "cannot append standalone comments to a populated line"
1478                 )
1479
1480         self.append(leaf, preformatted=preformatted)
1481
1482     @property
1483     def is_comment(self) -> bool:
1484         """Is this line a standalone comment?"""
1485         return len(self.leaves) == 1 and self.leaves[0].type == STANDALONE_COMMENT
1486
1487     @property
1488     def is_decorator(self) -> bool:
1489         """Is this line a decorator?"""
1490         return bool(self) and self.leaves[0].type == token.AT
1491
1492     @property
1493     def is_import(self) -> bool:
1494         """Is this an import line?"""
1495         return bool(self) and is_import(self.leaves[0])
1496
1497     @property
1498     def is_class(self) -> bool:
1499         """Is this line a class definition?"""
1500         return (
1501             bool(self)
1502             and self.leaves[0].type == token.NAME
1503             and self.leaves[0].value == "class"
1504         )
1505
1506     @property
1507     def is_stub_class(self) -> bool:
1508         """Is this line a class definition with a body consisting only of "..."?"""
1509         return self.is_class and self.leaves[-3:] == [
1510             Leaf(token.DOT, ".") for _ in range(3)
1511         ]
1512
1513     @property
1514     def is_def(self) -> bool:
1515         """Is this a function definition? (Also returns True for async defs.)"""
1516         try:
1517             first_leaf = self.leaves[0]
1518         except IndexError:
1519             return False
1520
1521         try:
1522             second_leaf: Optional[Leaf] = self.leaves[1]
1523         except IndexError:
1524             second_leaf = None
1525         return (first_leaf.type == token.NAME and first_leaf.value == "def") or (
1526             first_leaf.type == token.ASYNC
1527             and second_leaf is not None
1528             and second_leaf.type == token.NAME
1529             and second_leaf.value == "def"
1530         )
1531
1532     @property
1533     def is_class_paren_empty(self) -> bool:
1534         """Is this a class with no base classes but using parentheses?
1535
1536         Those are unnecessary and should be removed.
1537         """
1538         return (
1539             bool(self)
1540             and len(self.leaves) == 4
1541             and self.is_class
1542             and self.leaves[2].type == token.LPAR
1543             and self.leaves[2].value == "("
1544             and self.leaves[3].type == token.RPAR
1545             and self.leaves[3].value == ")"
1546         )
1547
1548     @property
1549     def is_triple_quoted_string(self) -> bool:
1550         """Is the line a triple quoted string?"""
1551         return (
1552             bool(self)
1553             and self.leaves[0].type == token.STRING
1554             and self.leaves[0].value.startswith(('"""', "'''"))
1555         )
1556
1557     def contains_standalone_comments(self, depth_limit: int = sys.maxsize) -> bool:
1558         """If so, needs to be split before emitting."""
1559         for leaf in self.leaves:
1560             if leaf.type == STANDALONE_COMMENT and leaf.bracket_depth <= depth_limit:
1561                 return True
1562
1563         return False
1564
1565     def contains_uncollapsable_type_comments(self) -> bool:
1566         ignored_ids = set()
1567         try:
1568             last_leaf = self.leaves[-1]
1569             ignored_ids.add(id(last_leaf))
1570             if last_leaf.type == token.COMMA or (
1571                 last_leaf.type == token.RPAR and not last_leaf.value
1572             ):
1573                 # When trailing commas or optional parens are inserted by Black for
1574                 # consistency, comments after the previous last element are not moved
1575                 # (they don't have to, rendering will still be correct).  So we ignore
1576                 # trailing commas and invisible.
1577                 last_leaf = self.leaves[-2]
1578                 ignored_ids.add(id(last_leaf))
1579         except IndexError:
1580             return False
1581
1582         # A type comment is uncollapsable if it is attached to a leaf
1583         # that isn't at the end of the line (since that could cause it
1584         # to get associated to a different argument) or if there are
1585         # comments before it (since that could cause it to get hidden
1586         # behind a comment.
1587         comment_seen = False
1588         for leaf_id, comments in self.comments.items():
1589             for comment in comments:
1590                 if is_type_comment(comment):
1591                     if comment_seen or (
1592                         not is_type_comment(comment, " ignore")
1593                         and leaf_id not in ignored_ids
1594                     ):
1595                         return True
1596
1597                 comment_seen = True
1598
1599         return False
1600
1601     def contains_unsplittable_type_ignore(self) -> bool:
1602         if not self.leaves:
1603             return False
1604
1605         # If a 'type: ignore' is attached to the end of a line, we
1606         # can't split the line, because we can't know which of the
1607         # subexpressions the ignore was meant to apply to.
1608         #
1609         # We only want this to apply to actual physical lines from the
1610         # original source, though: we don't want the presence of a
1611         # 'type: ignore' at the end of a multiline expression to
1612         # justify pushing it all onto one line. Thus we
1613         # (unfortunately) need to check the actual source lines and
1614         # only report an unsplittable 'type: ignore' if this line was
1615         # one line in the original code.
1616
1617         # Grab the first and last line numbers, skipping generated leaves
1618         first_line = next((leaf.lineno for leaf in self.leaves if leaf.lineno != 0), 0)
1619         last_line = next(
1620             (leaf.lineno for leaf in reversed(self.leaves) if leaf.lineno != 0), 0
1621         )
1622
1623         if first_line == last_line:
1624             # We look at the last two leaves since a comma or an
1625             # invisible paren could have been added at the end of the
1626             # line.
1627             for node in self.leaves[-2:]:
1628                 for comment in self.comments.get(id(node), []):
1629                     if is_type_comment(comment, " ignore"):
1630                         return True
1631
1632         return False
1633
1634     def contains_multiline_strings(self) -> bool:
1635         return any(is_multiline_string(leaf) for leaf in self.leaves)
1636
1637     def maybe_should_explode(self, closing: Leaf) -> bool:
1638         """Return True if this line should explode (always be split), that is when:
1639         - there's a trailing comma here; and
1640         - it's not a one-tuple.
1641         """
1642         if not (
1643             closing.type in CLOSING_BRACKETS
1644             and self.leaves
1645             and self.leaves[-1].type == token.COMMA
1646         ):
1647             return False
1648
1649         if closing.type in {token.RBRACE, token.RSQB}:
1650             return True
1651
1652         if self.is_import:
1653             return True
1654
1655         if not is_one_tuple_between(closing.opening_bracket, closing, self.leaves):
1656             return True
1657
1658         return False
1659
1660     def append_comment(self, comment: Leaf) -> bool:
1661         """Add an inline or standalone comment to the line."""
1662         if (
1663             comment.type == STANDALONE_COMMENT
1664             and self.bracket_tracker.any_open_brackets()
1665         ):
1666             comment.prefix = ""
1667             return False
1668
1669         if comment.type != token.COMMENT:
1670             return False
1671
1672         if not self.leaves:
1673             comment.type = STANDALONE_COMMENT
1674             comment.prefix = ""
1675             return False
1676
1677         last_leaf = self.leaves[-1]
1678         if (
1679             last_leaf.type == token.RPAR
1680             and not last_leaf.value
1681             and last_leaf.parent
1682             and len(list(last_leaf.parent.leaves())) <= 3
1683             and not is_type_comment(comment)
1684         ):
1685             # Comments on an optional parens wrapping a single leaf should belong to
1686             # the wrapped node except if it's a type comment. Pinning the comment like
1687             # this avoids unstable formatting caused by comment migration.
1688             if len(self.leaves) < 2:
1689                 comment.type = STANDALONE_COMMENT
1690                 comment.prefix = ""
1691                 return False
1692
1693             last_leaf = self.leaves[-2]
1694         self.comments.setdefault(id(last_leaf), []).append(comment)
1695         return True
1696
1697     def comments_after(self, leaf: Leaf) -> List[Leaf]:
1698         """Generate comments that should appear directly after `leaf`."""
1699         return self.comments.get(id(leaf), [])
1700
1701     def remove_trailing_comma(self) -> None:
1702         """Remove the trailing comma and moves the comments attached to it."""
1703         trailing_comma = self.leaves.pop()
1704         trailing_comma_comments = self.comments.pop(id(trailing_comma), [])
1705         self.comments.setdefault(id(self.leaves[-1]), []).extend(
1706             trailing_comma_comments
1707         )
1708
1709     def is_complex_subscript(self, leaf: Leaf) -> bool:
1710         """Return True iff `leaf` is part of a slice with non-trivial exprs."""
1711         open_lsqb = self.bracket_tracker.get_open_lsqb()
1712         if open_lsqb is None:
1713             return False
1714
1715         subscript_start = open_lsqb.next_sibling
1716
1717         if isinstance(subscript_start, Node):
1718             if subscript_start.type == syms.listmaker:
1719                 return False
1720
1721             if subscript_start.type == syms.subscriptlist:
1722                 subscript_start = child_towards(subscript_start, leaf)
1723         return subscript_start is not None and any(
1724             n.type in TEST_DESCENDANTS for n in subscript_start.pre_order()
1725         )
1726
1727     def clone(self) -> "Line":
1728         return Line(
1729             depth=self.depth,
1730             inside_brackets=self.inside_brackets,
1731             should_explode=self.should_explode,
1732         )
1733
1734     def __str__(self) -> str:
1735         """Render the line."""
1736         if not self:
1737             return "\n"
1738
1739         indent = "    " * self.depth
1740         leaves = iter(self.leaves)
1741         first = next(leaves)
1742         res = f"{first.prefix}{indent}{first.value}"
1743         for leaf in leaves:
1744             res += str(leaf)
1745         for comment in itertools.chain.from_iterable(self.comments.values()):
1746             res += str(comment)
1747
1748         return res + "\n"
1749
1750     def __bool__(self) -> bool:
1751         """Return True if the line has leaves or comments."""
1752         return bool(self.leaves or self.comments)
1753
1754
1755 @dataclass
1756 class EmptyLineTracker:
1757     """Provides a stateful method that returns the number of potential extra
1758     empty lines needed before and after the currently processed line.
1759
1760     Note: this tracker works on lines that haven't been split yet.  It assumes
1761     the prefix of the first leaf consists of optional newlines.  Those newlines
1762     are consumed by `maybe_empty_lines()` and included in the computation.
1763     """
1764
1765     is_pyi: bool = False
1766     previous_line: Optional[Line] = None
1767     previous_after: int = 0
1768     previous_defs: List[int] = field(default_factory=list)
1769
1770     def maybe_empty_lines(self, current_line: Line) -> Tuple[int, int]:
1771         """Return the number of extra empty lines before and after the `current_line`.
1772
1773         This is for separating `def`, `async def` and `class` with extra empty
1774         lines (two on module-level).
1775         """
1776         before, after = self._maybe_empty_lines(current_line)
1777         before = (
1778             # Black should not insert empty lines at the beginning
1779             # of the file
1780             0
1781             if self.previous_line is None
1782             else before - self.previous_after
1783         )
1784         self.previous_after = after
1785         self.previous_line = current_line
1786         return before, after
1787
1788     def _maybe_empty_lines(self, current_line: Line) -> Tuple[int, int]:
1789         max_allowed = 1
1790         if current_line.depth == 0:
1791             max_allowed = 1 if self.is_pyi else 2
1792         if current_line.leaves:
1793             # Consume the first leaf's extra newlines.
1794             first_leaf = current_line.leaves[0]
1795             before = first_leaf.prefix.count("\n")
1796             before = min(before, max_allowed)
1797             first_leaf.prefix = ""
1798         else:
1799             before = 0
1800         depth = current_line.depth
1801         while self.previous_defs and self.previous_defs[-1] >= depth:
1802             self.previous_defs.pop()
1803             if self.is_pyi:
1804                 before = 0 if depth else 1
1805             else:
1806                 before = 1 if depth else 2
1807         if current_line.is_decorator or current_line.is_def or current_line.is_class:
1808             return self._maybe_empty_lines_for_class_or_def(current_line, before)
1809
1810         if (
1811             self.previous_line
1812             and self.previous_line.is_import
1813             and not current_line.is_import
1814             and depth == self.previous_line.depth
1815         ):
1816             return (before or 1), 0
1817
1818         if (
1819             self.previous_line
1820             and self.previous_line.is_class
1821             and current_line.is_triple_quoted_string
1822         ):
1823             return before, 1
1824
1825         return before, 0
1826
1827     def _maybe_empty_lines_for_class_or_def(
1828         self, current_line: Line, before: int
1829     ) -> Tuple[int, int]:
1830         if not current_line.is_decorator:
1831             self.previous_defs.append(current_line.depth)
1832         if self.previous_line is None:
1833             # Don't insert empty lines before the first line in the file.
1834             return 0, 0
1835
1836         if self.previous_line.is_decorator:
1837             return 0, 0
1838
1839         if self.previous_line.depth < current_line.depth and (
1840             self.previous_line.is_class or self.previous_line.is_def
1841         ):
1842             return 0, 0
1843
1844         if (
1845             self.previous_line.is_comment
1846             and self.previous_line.depth == current_line.depth
1847             and before == 0
1848         ):
1849             return 0, 0
1850
1851         if self.is_pyi:
1852             if self.previous_line.depth > current_line.depth:
1853                 newlines = 1
1854             elif current_line.is_class or self.previous_line.is_class:
1855                 if current_line.is_stub_class and self.previous_line.is_stub_class:
1856                     # No blank line between classes with an empty body
1857                     newlines = 0
1858                 else:
1859                     newlines = 1
1860             elif current_line.is_def and not self.previous_line.is_def:
1861                 # Blank line between a block of functions and a block of non-functions
1862                 newlines = 1
1863             else:
1864                 newlines = 0
1865         else:
1866             newlines = 2
1867         if current_line.depth and newlines:
1868             newlines -= 1
1869         return newlines, 0
1870
1871
1872 @dataclass
1873 class LineGenerator(Visitor[Line]):
1874     """Generates reformatted Line objects.  Empty lines are not emitted.
1875
1876     Note: destroys the tree it's visiting by mutating prefixes of its leaves
1877     in ways that will no longer stringify to valid Python code on the tree.
1878     """
1879
1880     is_pyi: bool = False
1881     normalize_strings: bool = True
1882     current_line: Line = field(default_factory=Line)
1883     remove_u_prefix: bool = False
1884
1885     def line(self, indent: int = 0) -> Iterator[Line]:
1886         """Generate a line.
1887
1888         If the line is empty, only emit if it makes sense.
1889         If the line is too long, split it first and then generate.
1890
1891         If any lines were generated, set up a new current_line.
1892         """
1893         if not self.current_line:
1894             self.current_line.depth += indent
1895             return  # Line is empty, don't emit. Creating a new one unnecessary.
1896
1897         complete_line = self.current_line
1898         self.current_line = Line(depth=complete_line.depth + indent)
1899         yield complete_line
1900
1901     def visit_default(self, node: LN) -> Iterator[Line]:
1902         """Default `visit_*()` implementation. Recurses to children of `node`."""
1903         if isinstance(node, Leaf):
1904             any_open_brackets = self.current_line.bracket_tracker.any_open_brackets()
1905             for comment in generate_comments(node):
1906                 if any_open_brackets:
1907                     # any comment within brackets is subject to splitting
1908                     self.current_line.append(comment)
1909                 elif comment.type == token.COMMENT:
1910                     # regular trailing comment
1911                     self.current_line.append(comment)
1912                     yield from self.line()
1913
1914                 else:
1915                     # regular standalone comment
1916                     yield from self.line()
1917
1918                     self.current_line.append(comment)
1919                     yield from self.line()
1920
1921             normalize_prefix(node, inside_brackets=any_open_brackets)
1922             if self.normalize_strings and node.type == token.STRING:
1923                 normalize_string_prefix(node, remove_u_prefix=self.remove_u_prefix)
1924                 normalize_string_quotes(node)
1925             if node.type == token.NUMBER:
1926                 normalize_numeric_literal(node)
1927             if node.type not in WHITESPACE:
1928                 self.current_line.append(node)
1929         yield from super().visit_default(node)
1930
1931     def visit_INDENT(self, node: Leaf) -> Iterator[Line]:
1932         """Increase indentation level, maybe yield a line."""
1933         # In blib2to3 INDENT never holds comments.
1934         yield from self.line(+1)
1935         yield from self.visit_default(node)
1936
1937     def visit_DEDENT(self, node: Leaf) -> Iterator[Line]:
1938         """Decrease indentation level, maybe yield a line."""
1939         # The current line might still wait for trailing comments.  At DEDENT time
1940         # there won't be any (they would be prefixes on the preceding NEWLINE).
1941         # Emit the line then.
1942         yield from self.line()
1943
1944         # While DEDENT has no value, its prefix may contain standalone comments
1945         # that belong to the current indentation level.  Get 'em.
1946         yield from self.visit_default(node)
1947
1948         # Finally, emit the dedent.
1949         yield from self.line(-1)
1950
1951     def visit_stmt(
1952         self, node: Node, keywords: Set[str], parens: Set[str]
1953     ) -> Iterator[Line]:
1954         """Visit a statement.
1955
1956         This implementation is shared for `if`, `while`, `for`, `try`, `except`,
1957         `def`, `with`, `class`, `assert` and assignments.
1958
1959         The relevant Python language `keywords` for a given statement will be
1960         NAME leaves within it. This methods puts those on a separate line.
1961
1962         `parens` holds a set of string leaf values immediately after which
1963         invisible parens should be put.
1964         """
1965         normalize_invisible_parens(node, parens_after=parens)
1966         for child in node.children:
1967             if child.type == token.NAME and child.value in keywords:  # type: ignore
1968                 yield from self.line()
1969
1970             yield from self.visit(child)
1971
1972     def visit_suite(self, node: Node) -> Iterator[Line]:
1973         """Visit a suite."""
1974         if self.is_pyi and is_stub_suite(node):
1975             yield from self.visit(node.children[2])
1976         else:
1977             yield from self.visit_default(node)
1978
1979     def visit_simple_stmt(self, node: Node) -> Iterator[Line]:
1980         """Visit a statement without nested statements."""
1981         is_suite_like = node.parent and node.parent.type in STATEMENT
1982         if is_suite_like:
1983             if self.is_pyi and is_stub_body(node):
1984                 yield from self.visit_default(node)
1985             else:
1986                 yield from self.line(+1)
1987                 yield from self.visit_default(node)
1988                 yield from self.line(-1)
1989
1990         else:
1991             if not self.is_pyi or not node.parent or not is_stub_suite(node.parent):
1992                 yield from self.line()
1993             yield from self.visit_default(node)
1994
1995     def visit_async_stmt(self, node: Node) -> Iterator[Line]:
1996         """Visit `async def`, `async for`, `async with`."""
1997         yield from self.line()
1998
1999         children = iter(node.children)
2000         for child in children:
2001             yield from self.visit(child)
2002
2003             if child.type == token.ASYNC:
2004                 break
2005
2006         internal_stmt = next(children)
2007         for child in internal_stmt.children:
2008             yield from self.visit(child)
2009
2010     def visit_decorators(self, node: Node) -> Iterator[Line]:
2011         """Visit decorators."""
2012         for child in node.children:
2013             yield from self.line()
2014             yield from self.visit(child)
2015
2016     def visit_SEMI(self, leaf: Leaf) -> Iterator[Line]:
2017         """Remove a semicolon and put the other statement on a separate line."""
2018         yield from self.line()
2019
2020     def visit_ENDMARKER(self, leaf: Leaf) -> Iterator[Line]:
2021         """End of file. Process outstanding comments and end with a newline."""
2022         yield from self.visit_default(leaf)
2023         yield from self.line()
2024
2025     def visit_STANDALONE_COMMENT(self, leaf: Leaf) -> Iterator[Line]:
2026         if not self.current_line.bracket_tracker.any_open_brackets():
2027             yield from self.line()
2028         yield from self.visit_default(leaf)
2029
2030     def visit_factor(self, node: Node) -> Iterator[Line]:
2031         """Force parentheses between a unary op and a binary power:
2032
2033         -2 ** 8 -> -(2 ** 8)
2034         """
2035         _operator, operand = node.children
2036         if (
2037             operand.type == syms.power
2038             and len(operand.children) == 3
2039             and operand.children[1].type == token.DOUBLESTAR
2040         ):
2041             lpar = Leaf(token.LPAR, "(")
2042             rpar = Leaf(token.RPAR, ")")
2043             index = operand.remove() or 0
2044             node.insert_child(index, Node(syms.atom, [lpar, operand, rpar]))
2045         yield from self.visit_default(node)
2046
2047     def visit_STRING(self, leaf: Leaf) -> Iterator[Line]:
2048         if is_docstring(leaf) and "\\\n" not in leaf.value:
2049             # We're ignoring docstrings with backslash newline escapes because changing
2050             # indentation of those changes the AST representation of the code.
2051             prefix = get_string_prefix(leaf.value)
2052             lead_len = len(prefix) + 3
2053             tail_len = -3
2054             indent = " " * 4 * self.current_line.depth
2055             docstring = fix_docstring(leaf.value[lead_len:tail_len], indent)
2056             if docstring:
2057                 if leaf.value[lead_len - 1] == docstring[0]:
2058                     docstring = " " + docstring
2059                 if leaf.value[tail_len + 1] == docstring[-1]:
2060                     docstring = docstring + " "
2061             leaf.value = leaf.value[0:lead_len] + docstring + leaf.value[tail_len:]
2062
2063         yield from self.visit_default(leaf)
2064
2065     def __post_init__(self) -> None:
2066         """You are in a twisty little maze of passages."""
2067         v = self.visit_stmt
2068         Ø: Set[str] = set()
2069         self.visit_assert_stmt = partial(v, keywords={"assert"}, parens={"assert", ","})
2070         self.visit_if_stmt = partial(
2071             v, keywords={"if", "else", "elif"}, parens={"if", "elif"}
2072         )
2073         self.visit_while_stmt = partial(v, keywords={"while", "else"}, parens={"while"})
2074         self.visit_for_stmt = partial(v, keywords={"for", "else"}, parens={"for", "in"})
2075         self.visit_try_stmt = partial(
2076             v, keywords={"try", "except", "else", "finally"}, parens=Ø
2077         )
2078         self.visit_except_clause = partial(v, keywords={"except"}, parens=Ø)
2079         self.visit_with_stmt = partial(v, keywords={"with"}, parens=Ø)
2080         self.visit_funcdef = partial(v, keywords={"def"}, parens=Ø)
2081         self.visit_classdef = partial(v, keywords={"class"}, parens=Ø)
2082         self.visit_expr_stmt = partial(v, keywords=Ø, parens=ASSIGNMENTS)
2083         self.visit_return_stmt = partial(v, keywords={"return"}, parens={"return"})
2084         self.visit_import_from = partial(v, keywords=Ø, parens={"import"})
2085         self.visit_del_stmt = partial(v, keywords=Ø, parens={"del"})
2086         self.visit_async_funcdef = self.visit_async_stmt
2087         self.visit_decorated = self.visit_decorators
2088
2089
2090 IMPLICIT_TUPLE = {syms.testlist, syms.testlist_star_expr, syms.exprlist}
2091 BRACKET = {token.LPAR: token.RPAR, token.LSQB: token.RSQB, token.LBRACE: token.RBRACE}
2092 OPENING_BRACKETS = set(BRACKET.keys())
2093 CLOSING_BRACKETS = set(BRACKET.values())
2094 BRACKETS = OPENING_BRACKETS | CLOSING_BRACKETS
2095 ALWAYS_NO_SPACE = CLOSING_BRACKETS | {token.COMMA, STANDALONE_COMMENT}
2096
2097
2098 def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str:  # noqa: C901
2099     """Return whitespace prefix if needed for the given `leaf`.
2100
2101     `complex_subscript` signals whether the given leaf is part of a subscription
2102     which has non-trivial arguments, like arithmetic expressions or function calls.
2103     """
2104     NO = ""
2105     SPACE = " "
2106     DOUBLESPACE = "  "
2107     t = leaf.type
2108     p = leaf.parent
2109     v = leaf.value
2110     if t in ALWAYS_NO_SPACE:
2111         return NO
2112
2113     if t == token.COMMENT:
2114         return DOUBLESPACE
2115
2116     assert p is not None, f"INTERNAL ERROR: hand-made leaf without parent: {leaf!r}"
2117     if t == token.COLON and p.type not in {
2118         syms.subscript,
2119         syms.subscriptlist,
2120         syms.sliceop,
2121     }:
2122         return NO
2123
2124     prev = leaf.prev_sibling
2125     if not prev:
2126         prevp = preceding_leaf(p)
2127         if not prevp or prevp.type in OPENING_BRACKETS:
2128             return NO
2129
2130         if t == token.COLON:
2131             if prevp.type == token.COLON:
2132                 return NO
2133
2134             elif prevp.type != token.COMMA and not complex_subscript:
2135                 return NO
2136
2137             return SPACE
2138
2139         if prevp.type == token.EQUAL:
2140             if prevp.parent:
2141                 if prevp.parent.type in {
2142                     syms.arglist,
2143                     syms.argument,
2144                     syms.parameters,
2145                     syms.varargslist,
2146                 }:
2147                     return NO
2148
2149                 elif prevp.parent.type == syms.typedargslist:
2150                     # A bit hacky: if the equal sign has whitespace, it means we
2151                     # previously found it's a typed argument.  So, we're using
2152                     # that, too.
2153                     return prevp.prefix
2154
2155         elif prevp.type in VARARGS_SPECIALS:
2156             if is_vararg(prevp, within=VARARGS_PARENTS | UNPACKING_PARENTS):
2157                 return NO
2158
2159         elif prevp.type == token.COLON:
2160             if prevp.parent and prevp.parent.type in {syms.subscript, syms.sliceop}:
2161                 return SPACE if complex_subscript else NO
2162
2163         elif (
2164             prevp.parent
2165             and prevp.parent.type == syms.factor
2166             and prevp.type in MATH_OPERATORS
2167         ):
2168             return NO
2169
2170         elif (
2171             prevp.type == token.RIGHTSHIFT
2172             and prevp.parent
2173             and prevp.parent.type == syms.shift_expr
2174             and prevp.prev_sibling
2175             and prevp.prev_sibling.type == token.NAME
2176             and prevp.prev_sibling.value == "print"  # type: ignore
2177         ):
2178             # Python 2 print chevron
2179             return NO
2180
2181     elif prev.type in OPENING_BRACKETS:
2182         return NO
2183
2184     if p.type in {syms.parameters, syms.arglist}:
2185         # untyped function signatures or calls
2186         if not prev or prev.type != token.COMMA:
2187             return NO
2188
2189     elif p.type == syms.varargslist:
2190         # lambdas
2191         if prev and prev.type != token.COMMA:
2192             return NO
2193
2194     elif p.type == syms.typedargslist:
2195         # typed function signatures
2196         if not prev:
2197             return NO
2198
2199         if t == token.EQUAL:
2200             if prev.type != syms.tname:
2201                 return NO
2202
2203         elif prev.type == token.EQUAL:
2204             # A bit hacky: if the equal sign has whitespace, it means we
2205             # previously found it's a typed argument.  So, we're using that, too.
2206             return prev.prefix
2207
2208         elif prev.type != token.COMMA:
2209             return NO
2210
2211     elif p.type == syms.tname:
2212         # type names
2213         if not prev:
2214             prevp = preceding_leaf(p)
2215             if not prevp or prevp.type != token.COMMA:
2216                 return NO
2217
2218     elif p.type == syms.trailer:
2219         # attributes and calls
2220         if t == token.LPAR or t == token.RPAR:
2221             return NO
2222
2223         if not prev:
2224             if t == token.DOT:
2225                 prevp = preceding_leaf(p)
2226                 if not prevp or prevp.type != token.NUMBER:
2227                     return NO
2228
2229             elif t == token.LSQB:
2230                 return NO
2231
2232         elif prev.type != token.COMMA:
2233             return NO
2234
2235     elif p.type == syms.argument:
2236         # single argument
2237         if t == token.EQUAL:
2238             return NO
2239
2240         if not prev:
2241             prevp = preceding_leaf(p)
2242             if not prevp or prevp.type == token.LPAR:
2243                 return NO
2244
2245         elif prev.type in {token.EQUAL} | VARARGS_SPECIALS:
2246             return NO
2247
2248     elif p.type == syms.decorator:
2249         # decorators
2250         return NO
2251
2252     elif p.type == syms.dotted_name:
2253         if prev:
2254             return NO
2255
2256         prevp = preceding_leaf(p)
2257         if not prevp or prevp.type == token.AT or prevp.type == token.DOT:
2258             return NO
2259
2260     elif p.type == syms.classdef:
2261         if t == token.LPAR:
2262             return NO
2263
2264         if prev and prev.type == token.LPAR:
2265             return NO
2266
2267     elif p.type in {syms.subscript, syms.sliceop}:
2268         # indexing
2269         if not prev:
2270             assert p.parent is not None, "subscripts are always parented"
2271             if p.parent.type == syms.subscriptlist:
2272                 return SPACE
2273
2274             return NO
2275
2276         elif not complex_subscript:
2277             return NO
2278
2279     elif p.type == syms.atom:
2280         if prev and t == token.DOT:
2281             # dots, but not the first one.
2282             return NO
2283
2284     elif p.type == syms.dictsetmaker:
2285         # dict unpacking
2286         if prev and prev.type == token.DOUBLESTAR:
2287             return NO
2288
2289     elif p.type in {syms.factor, syms.star_expr}:
2290         # unary ops
2291         if not prev:
2292             prevp = preceding_leaf(p)
2293             if not prevp or prevp.type in OPENING_BRACKETS:
2294                 return NO
2295
2296             prevp_parent = prevp.parent
2297             assert prevp_parent is not None
2298             if prevp.type == token.COLON and prevp_parent.type in {
2299                 syms.subscript,
2300                 syms.sliceop,
2301             }:
2302                 return NO
2303
2304             elif prevp.type == token.EQUAL and prevp_parent.type == syms.argument:
2305                 return NO
2306
2307         elif t in {token.NAME, token.NUMBER, token.STRING}:
2308             return NO
2309
2310     elif p.type == syms.import_from:
2311         if t == token.DOT:
2312             if prev and prev.type == token.DOT:
2313                 return NO
2314
2315         elif t == token.NAME:
2316             if v == "import":
2317                 return SPACE
2318
2319             if prev and prev.type == token.DOT:
2320                 return NO
2321
2322     elif p.type == syms.sliceop:
2323         return NO
2324
2325     return SPACE
2326
2327
2328 def preceding_leaf(node: Optional[LN]) -> Optional[Leaf]:
2329     """Return the first leaf that precedes `node`, if any."""
2330     while node:
2331         res = node.prev_sibling
2332         if res:
2333             if isinstance(res, Leaf):
2334                 return res
2335
2336             try:
2337                 return list(res.leaves())[-1]
2338
2339             except IndexError:
2340                 return None
2341
2342         node = node.parent
2343     return None
2344
2345
2346 def prev_siblings_are(node: Optional[LN], tokens: List[Optional[NodeType]]) -> bool:
2347     """Return if the `node` and its previous siblings match types against the provided
2348     list of tokens; the provided `node`has its type matched against the last element in
2349     the list.  `None` can be used as the first element to declare that the start of the
2350     list is anchored at the start of its parent's children."""
2351     if not tokens:
2352         return True
2353     if tokens[-1] is None:
2354         return node is None
2355     if not node:
2356         return False
2357     if node.type != tokens[-1]:
2358         return False
2359     return prev_siblings_are(node.prev_sibling, tokens[:-1])
2360
2361
2362 def child_towards(ancestor: Node, descendant: LN) -> Optional[LN]:
2363     """Return the child of `ancestor` that contains `descendant`."""
2364     node: Optional[LN] = descendant
2365     while node and node.parent != ancestor:
2366         node = node.parent
2367     return node
2368
2369
2370 def container_of(leaf: Leaf) -> LN:
2371     """Return `leaf` or one of its ancestors that is the topmost container of it.
2372
2373     By "container" we mean a node where `leaf` is the very first child.
2374     """
2375     same_prefix = leaf.prefix
2376     container: LN = leaf
2377     while container:
2378         parent = container.parent
2379         if parent is None:
2380             break
2381
2382         if parent.children[0].prefix != same_prefix:
2383             break
2384
2385         if parent.type == syms.file_input:
2386             break
2387
2388         if parent.prev_sibling is not None and parent.prev_sibling.type in BRACKETS:
2389             break
2390
2391         container = parent
2392     return container
2393
2394
2395 def is_split_after_delimiter(leaf: Leaf, previous: Optional[Leaf] = None) -> Priority:
2396     """Return the priority of the `leaf` delimiter, given a line break after it.
2397
2398     The delimiter priorities returned here are from those delimiters that would
2399     cause a line break after themselves.
2400
2401     Higher numbers are higher priority.
2402     """
2403     if leaf.type == token.COMMA:
2404         return COMMA_PRIORITY
2405
2406     return 0
2407
2408
2409 def is_split_before_delimiter(leaf: Leaf, previous: Optional[Leaf] = None) -> Priority:
2410     """Return the priority of the `leaf` delimiter, given a line break before it.
2411
2412     The delimiter priorities returned here are from those delimiters that would
2413     cause a line break before themselves.
2414
2415     Higher numbers are higher priority.
2416     """
2417     if is_vararg(leaf, within=VARARGS_PARENTS | UNPACKING_PARENTS):
2418         # * and ** might also be MATH_OPERATORS but in this case they are not.
2419         # Don't treat them as a delimiter.
2420         return 0
2421
2422     if (
2423         leaf.type == token.DOT
2424         and leaf.parent
2425         and leaf.parent.type not in {syms.import_from, syms.dotted_name}
2426         and (previous is None or previous.type in CLOSING_BRACKETS)
2427     ):
2428         return DOT_PRIORITY
2429
2430     if (
2431         leaf.type in MATH_OPERATORS
2432         and leaf.parent
2433         and leaf.parent.type not in {syms.factor, syms.star_expr}
2434     ):
2435         return MATH_PRIORITIES[leaf.type]
2436
2437     if leaf.type in COMPARATORS:
2438         return COMPARATOR_PRIORITY
2439
2440     if (
2441         leaf.type == token.STRING
2442         and previous is not None
2443         and previous.type == token.STRING
2444     ):
2445         return STRING_PRIORITY
2446
2447     if leaf.type not in {token.NAME, token.ASYNC}:
2448         return 0
2449
2450     if (
2451         leaf.value == "for"
2452         and leaf.parent
2453         and leaf.parent.type in {syms.comp_for, syms.old_comp_for}
2454         or leaf.type == token.ASYNC
2455     ):
2456         if (
2457             not isinstance(leaf.prev_sibling, Leaf)
2458             or leaf.prev_sibling.value != "async"
2459         ):
2460             return COMPREHENSION_PRIORITY
2461
2462     if (
2463         leaf.value == "if"
2464         and leaf.parent
2465         and leaf.parent.type in {syms.comp_if, syms.old_comp_if}
2466     ):
2467         return COMPREHENSION_PRIORITY
2468
2469     if leaf.value in {"if", "else"} and leaf.parent and leaf.parent.type == syms.test:
2470         return TERNARY_PRIORITY
2471
2472     if leaf.value == "is":
2473         return COMPARATOR_PRIORITY
2474
2475     if (
2476         leaf.value == "in"
2477         and leaf.parent
2478         and leaf.parent.type in {syms.comp_op, syms.comparison}
2479         and not (
2480             previous is not None
2481             and previous.type == token.NAME
2482             and previous.value == "not"
2483         )
2484     ):
2485         return COMPARATOR_PRIORITY
2486
2487     if (
2488         leaf.value == "not"
2489         and leaf.parent
2490         and leaf.parent.type == syms.comp_op
2491         and not (
2492             previous is not None
2493             and previous.type == token.NAME
2494             and previous.value == "is"
2495         )
2496     ):
2497         return COMPARATOR_PRIORITY
2498
2499     if leaf.value in LOGIC_OPERATORS and leaf.parent:
2500         return LOGIC_PRIORITY
2501
2502     return 0
2503
2504
2505 FMT_OFF = {"# fmt: off", "# fmt:off", "# yapf: disable"}
2506 FMT_ON = {"# fmt: on", "# fmt:on", "# yapf: enable"}
2507
2508
2509 def generate_comments(leaf: LN) -> Iterator[Leaf]:
2510     """Clean the prefix of the `leaf` and generate comments from it, if any.
2511
2512     Comments in lib2to3 are shoved into the whitespace prefix.  This happens
2513     in `pgen2/driver.py:Driver.parse_tokens()`.  This was a brilliant implementation
2514     move because it does away with modifying the grammar to include all the
2515     possible places in which comments can be placed.
2516
2517     The sad consequence for us though is that comments don't "belong" anywhere.
2518     This is why this function generates simple parentless Leaf objects for
2519     comments.  We simply don't know what the correct parent should be.
2520
2521     No matter though, we can live without this.  We really only need to
2522     differentiate between inline and standalone comments.  The latter don't
2523     share the line with any code.
2524
2525     Inline comments are emitted as regular token.COMMENT leaves.  Standalone
2526     are emitted with a fake STANDALONE_COMMENT token identifier.
2527     """
2528     for pc in list_comments(leaf.prefix, is_endmarker=leaf.type == token.ENDMARKER):
2529         yield Leaf(pc.type, pc.value, prefix="\n" * pc.newlines)
2530
2531
2532 @dataclass
2533 class ProtoComment:
2534     """Describes a piece of syntax that is a comment.
2535
2536     It's not a :class:`blib2to3.pytree.Leaf` so that:
2537
2538     * it can be cached (`Leaf` objects should not be reused more than once as
2539       they store their lineno, column, prefix, and parent information);
2540     * `newlines` and `consumed` fields are kept separate from the `value`. This
2541       simplifies handling of special marker comments like ``# fmt: off/on``.
2542     """
2543
2544     type: int  # token.COMMENT or STANDALONE_COMMENT
2545     value: str  # content of the comment
2546     newlines: int  # how many newlines before the comment
2547     consumed: int  # how many characters of the original leaf's prefix did we consume
2548
2549
2550 @lru_cache(maxsize=4096)
2551 def list_comments(prefix: str, *, is_endmarker: bool) -> List[ProtoComment]:
2552     """Return a list of :class:`ProtoComment` objects parsed from the given `prefix`."""
2553     result: List[ProtoComment] = []
2554     if not prefix or "#" not in prefix:
2555         return result
2556
2557     consumed = 0
2558     nlines = 0
2559     ignored_lines = 0
2560     for index, line in enumerate(prefix.split("\n")):
2561         consumed += len(line) + 1  # adding the length of the split '\n'
2562         line = line.lstrip()
2563         if not line:
2564             nlines += 1
2565         if not line.startswith("#"):
2566             # Escaped newlines outside of a comment are not really newlines at
2567             # all. We treat a single-line comment following an escaped newline
2568             # as a simple trailing comment.
2569             if line.endswith("\\"):
2570                 ignored_lines += 1
2571             continue
2572
2573         if index == ignored_lines and not is_endmarker:
2574             comment_type = token.COMMENT  # simple trailing comment
2575         else:
2576             comment_type = STANDALONE_COMMENT
2577         comment = make_comment(line)
2578         result.append(
2579             ProtoComment(
2580                 type=comment_type, value=comment, newlines=nlines, consumed=consumed
2581             )
2582         )
2583         nlines = 0
2584     return result
2585
2586
2587 def make_comment(content: str) -> str:
2588     """Return a consistently formatted comment from the given `content` string.
2589
2590     All comments (except for "##", "#!", "#:", '#'", "#%%") should have a single
2591     space between the hash sign and the content.
2592
2593     If `content` didn't start with a hash sign, one is provided.
2594     """
2595     content = content.rstrip()
2596     if not content:
2597         return "#"
2598
2599     if content[0] == "#":
2600         content = content[1:]
2601     if content and content[0] not in " !:#'%":
2602         content = " " + content
2603     return "#" + content
2604
2605
2606 def transform_line(
2607     line: Line, mode: Mode, features: Collection[Feature] = ()
2608 ) -> Iterator[Line]:
2609     """Transform a `line`, potentially splitting it into many lines.
2610
2611     They should fit in the allotted `line_length` but might not be able to.
2612
2613     `features` are syntactical features that may be used in the output.
2614     """
2615     if line.is_comment:
2616         yield line
2617         return
2618
2619     line_str = line_to_string(line)
2620
2621     def init_st(ST: Type[StringTransformer]) -> StringTransformer:
2622         """Initialize StringTransformer"""
2623         return ST(mode.line_length, mode.string_normalization)
2624
2625     string_merge = init_st(StringMerger)
2626     string_paren_strip = init_st(StringParenStripper)
2627     string_split = init_st(StringSplitter)
2628     string_paren_wrap = init_st(StringParenWrapper)
2629
2630     transformers: List[Transformer]
2631     if (
2632         not line.contains_uncollapsable_type_comments()
2633         and not line.should_explode
2634         and (
2635             is_line_short_enough(line, line_length=mode.line_length, line_str=line_str)
2636             or line.contains_unsplittable_type_ignore()
2637         )
2638         and not (line.inside_brackets and line.contains_standalone_comments())
2639     ):
2640         # Only apply basic string preprocessing, since lines shouldn't be split here.
2641         if mode.experimental_string_processing:
2642             transformers = [string_merge, string_paren_strip]
2643         else:
2644             transformers = []
2645     elif line.is_def:
2646         transformers = [left_hand_split]
2647     else:
2648
2649         def rhs(line: Line, features: Collection[Feature]) -> Iterator[Line]:
2650             """Wraps calls to `right_hand_split`.
2651
2652             The calls increasingly `omit` right-hand trailers (bracket pairs with
2653             content), meaning the trailers get glued together to split on another
2654             bracket pair instead.
2655             """
2656             for omit in generate_trailers_to_omit(line, mode.line_length):
2657                 lines = list(
2658                     right_hand_split(line, mode.line_length, features, omit=omit)
2659                 )
2660                 # Note: this check is only able to figure out if the first line of the
2661                 # *current* transformation fits in the line length.  This is true only
2662                 # for simple cases.  All others require running more transforms via
2663                 # `transform_line()`.  This check doesn't know if those would succeed.
2664                 if is_line_short_enough(lines[0], line_length=mode.line_length):
2665                     yield from lines
2666                     return
2667
2668             # All splits failed, best effort split with no omits.
2669             # This mostly happens to multiline strings that are by definition
2670             # reported as not fitting a single line, as well as lines that contain
2671             # trailing commas (those have to be exploded).
2672             yield from right_hand_split(
2673                 line, line_length=mode.line_length, features=features
2674             )
2675
2676         if mode.experimental_string_processing:
2677             if line.inside_brackets:
2678                 transformers = [
2679                     string_merge,
2680                     string_paren_strip,
2681                     string_split,
2682                     delimiter_split,
2683                     standalone_comment_split,
2684                     string_paren_wrap,
2685                     rhs,
2686                 ]
2687             else:
2688                 transformers = [
2689                     string_merge,
2690                     string_paren_strip,
2691                     string_split,
2692                     string_paren_wrap,
2693                     rhs,
2694                 ]
2695         else:
2696             if line.inside_brackets:
2697                 transformers = [delimiter_split, standalone_comment_split, rhs]
2698             else:
2699                 transformers = [rhs]
2700
2701     for transform in transformers:
2702         # We are accumulating lines in `result` because we might want to abort
2703         # mission and return the original line in the end, or attempt a different
2704         # split altogether.
2705         try:
2706             result = run_transformer(line, transform, mode, features, line_str=line_str)
2707         except CannotTransform:
2708             continue
2709         else:
2710             yield from result
2711             break
2712
2713     else:
2714         yield line
2715
2716
2717 @dataclass  # type: ignore
2718 class StringTransformer(ABC):
2719     """
2720     An implementation of the Transformer protocol that relies on its
2721     subclasses overriding the template methods `do_match(...)` and
2722     `do_transform(...)`.
2723
2724     This Transformer works exclusively on strings (for example, by merging
2725     or splitting them).
2726
2727     The following sections can be found among the docstrings of each concrete
2728     StringTransformer subclass.
2729
2730     Requirements:
2731         Which requirements must be met of the given Line for this
2732         StringTransformer to be applied?
2733
2734     Transformations:
2735         If the given Line meets all of the above requirements, which string
2736         transformations can you expect to be applied to it by this
2737         StringTransformer?
2738
2739     Collaborations:
2740         What contractual agreements does this StringTransformer have with other
2741         StringTransfomers? Such collaborations should be eliminated/minimized
2742         as much as possible.
2743     """
2744
2745     line_length: int
2746     normalize_strings: bool
2747     __name__ = "StringTransformer"
2748
2749     @abstractmethod
2750     def do_match(self, line: Line) -> TMatchResult:
2751         """
2752         Returns:
2753             * Ok(string_idx) such that `line.leaves[string_idx]` is our target
2754             string, if a match was able to be made.
2755                 OR
2756             * Err(CannotTransform), if a match was not able to be made.
2757         """
2758
2759     @abstractmethod
2760     def do_transform(self, line: Line, string_idx: int) -> Iterator[TResult[Line]]:
2761         """
2762         Yields:
2763             * Ok(new_line) where new_line is the new transformed line.
2764                 OR
2765             * Err(CannotTransform) if the transformation failed for some reason. The
2766             `do_match(...)` template method should usually be used to reject
2767             the form of the given Line, but in some cases it is difficult to
2768             know whether or not a Line meets the StringTransformer's
2769             requirements until the transformation is already midway.
2770
2771         Side Effects:
2772             This method should NOT mutate @line directly, but it MAY mutate the
2773             Line's underlying Node structure. (WARNING: If the underlying Node
2774             structure IS altered, then this method should NOT be allowed to
2775             yield an CannotTransform after that point.)
2776         """
2777
2778     def __call__(self, line: Line, _features: Collection[Feature]) -> Iterator[Line]:
2779         """
2780         StringTransformer instances have a call signature that mirrors that of
2781         the Transformer type.
2782
2783         Raises:
2784             CannotTransform(...) if the concrete StringTransformer class is unable
2785             to transform @line.
2786         """
2787         # Optimization to avoid calling `self.do_match(...)` when the line does
2788         # not contain any string.
2789         if not any(leaf.type == token.STRING for leaf in line.leaves):
2790             raise CannotTransform("There are no strings in this line.")
2791
2792         match_result = self.do_match(line)
2793
2794         if isinstance(match_result, Err):
2795             cant_transform = match_result.err()
2796             raise CannotTransform(
2797                 f"The string transformer {self.__class__.__name__} does not recognize"
2798                 " this line as one that it can transform."
2799             ) from cant_transform
2800
2801         string_idx = match_result.ok()
2802
2803         for line_result in self.do_transform(line, string_idx):
2804             if isinstance(line_result, Err):
2805                 cant_transform = line_result.err()
2806                 raise CannotTransform(
2807                     "StringTransformer failed while attempting to transform string."
2808                 ) from cant_transform
2809             line = line_result.ok()
2810             yield line
2811
2812
2813 @dataclass
2814 class CustomSplit:
2815     """A custom (i.e. manual) string split.
2816
2817     A single CustomSplit instance represents a single substring.
2818
2819     Examples:
2820         Consider the following string:
2821         ```
2822         "Hi there friend."
2823         " This is a custom"
2824         f" string {split}."
2825         ```
2826
2827         This string will correspond to the following three CustomSplit instances:
2828         ```
2829         CustomSplit(False, 16)
2830         CustomSplit(False, 17)
2831         CustomSplit(True, 16)
2832         ```
2833     """
2834
2835     has_prefix: bool
2836     break_idx: int
2837
2838
2839 class CustomSplitMapMixin:
2840     """
2841     This mixin class is used to map merged strings to a sequence of
2842     CustomSplits, which will then be used to re-split the strings iff none of
2843     the resultant substrings go over the configured max line length.
2844     """
2845
2846     _Key = Tuple[StringID, str]
2847     _CUSTOM_SPLIT_MAP: Dict[_Key, Tuple[CustomSplit, ...]] = defaultdict(tuple)
2848
2849     @staticmethod
2850     def _get_key(string: str) -> "CustomSplitMapMixin._Key":
2851         """
2852         Returns:
2853             A unique identifier that is used internally to map @string to a
2854             group of custom splits.
2855         """
2856         return (id(string), string)
2857
2858     def add_custom_splits(
2859         self, string: str, custom_splits: Iterable[CustomSplit]
2860     ) -> None:
2861         """Custom Split Map Setter Method
2862
2863         Side Effects:
2864             Adds a mapping from @string to the custom splits @custom_splits.
2865         """
2866         key = self._get_key(string)
2867         self._CUSTOM_SPLIT_MAP[key] = tuple(custom_splits)
2868
2869     def pop_custom_splits(self, string: str) -> List[CustomSplit]:
2870         """Custom Split Map Getter Method
2871
2872         Returns:
2873             * A list of the custom splits that are mapped to @string, if any
2874             exist.
2875                 OR
2876             * [], otherwise.
2877
2878         Side Effects:
2879             Deletes the mapping between @string and its associated custom
2880             splits (which are returned to the caller).
2881         """
2882         key = self._get_key(string)
2883
2884         custom_splits = self._CUSTOM_SPLIT_MAP[key]
2885         del self._CUSTOM_SPLIT_MAP[key]
2886
2887         return list(custom_splits)
2888
2889     def has_custom_splits(self, string: str) -> bool:
2890         """
2891         Returns:
2892             True iff @string is associated with a set of custom splits.
2893         """
2894         key = self._get_key(string)
2895         return key in self._CUSTOM_SPLIT_MAP
2896
2897
2898 class StringMerger(CustomSplitMapMixin, StringTransformer):
2899     """StringTransformer that merges strings together.
2900
2901     Requirements:
2902         (A) The line contains adjacent strings such that ALL of the validation checks
2903         listed in StringMerger.__validate_msg(...)'s docstring pass.
2904             OR
2905         (B) The line contains a string which uses line continuation backslashes.
2906
2907     Transformations:
2908         Depending on which of the two requirements above where met, either:
2909
2910         (A) The string group associated with the target string is merged.
2911             OR
2912         (B) All line-continuation backslashes are removed from the target string.
2913
2914     Collaborations:
2915         StringMerger provides custom split information to StringSplitter.
2916     """
2917
2918     def do_match(self, line: Line) -> TMatchResult:
2919         LL = line.leaves
2920
2921         is_valid_index = is_valid_index_factory(LL)
2922
2923         for (i, leaf) in enumerate(LL):
2924             if (
2925                 leaf.type == token.STRING
2926                 and is_valid_index(i + 1)
2927                 and LL[i + 1].type == token.STRING
2928             ):
2929                 return Ok(i)
2930
2931             if leaf.type == token.STRING and "\\\n" in leaf.value:
2932                 return Ok(i)
2933
2934         return TErr("This line has no strings that need merging.")
2935
2936     def do_transform(self, line: Line, string_idx: int) -> Iterator[TResult[Line]]:
2937         new_line = line
2938         rblc_result = self.__remove_backslash_line_continuation_chars(
2939             new_line, string_idx
2940         )
2941         if isinstance(rblc_result, Ok):
2942             new_line = rblc_result.ok()
2943
2944         msg_result = self.__merge_string_group(new_line, string_idx)
2945         if isinstance(msg_result, Ok):
2946             new_line = msg_result.ok()
2947
2948         if isinstance(rblc_result, Err) and isinstance(msg_result, Err):
2949             msg_cant_transform = msg_result.err()
2950             rblc_cant_transform = rblc_result.err()
2951             cant_transform = CannotTransform(
2952                 "StringMerger failed to merge any strings in this line."
2953             )
2954
2955             # Chain the errors together using `__cause__`.
2956             msg_cant_transform.__cause__ = rblc_cant_transform
2957             cant_transform.__cause__ = msg_cant_transform
2958
2959             yield Err(cant_transform)
2960         else:
2961             yield Ok(new_line)
2962
2963     @staticmethod
2964     def __remove_backslash_line_continuation_chars(
2965         line: Line, string_idx: int
2966     ) -> TResult[Line]:
2967         """
2968         Merge strings that were split across multiple lines using
2969         line-continuation backslashes.
2970
2971         Returns:
2972             Ok(new_line), if @line contains backslash line-continuation
2973             characters.
2974                 OR
2975             Err(CannotTransform), otherwise.
2976         """
2977         LL = line.leaves
2978
2979         string_leaf = LL[string_idx]
2980         if not (
2981             string_leaf.type == token.STRING
2982             and "\\\n" in string_leaf.value
2983             and not has_triple_quotes(string_leaf.value)
2984         ):
2985             return TErr(
2986                 f"String leaf {string_leaf} does not contain any backslash line"
2987                 " continuation characters."
2988             )
2989
2990         new_line = line.clone()
2991         new_line.comments = line.comments.copy()
2992         append_leaves(new_line, line, LL)
2993
2994         new_string_leaf = new_line.leaves[string_idx]
2995         new_string_leaf.value = new_string_leaf.value.replace("\\\n", "")
2996
2997         return Ok(new_line)
2998
2999     def __merge_string_group(self, line: Line, string_idx: int) -> TResult[Line]:
3000         """
3001         Merges string group (i.e. set of adjacent strings) where the first
3002         string in the group is `line.leaves[string_idx]`.
3003
3004         Returns:
3005             Ok(new_line), if ALL of the validation checks found in
3006             __validate_msg(...) pass.
3007                 OR
3008             Err(CannotTransform), otherwise.
3009         """
3010         LL = line.leaves
3011
3012         is_valid_index = is_valid_index_factory(LL)
3013
3014         vresult = self.__validate_msg(line, string_idx)
3015         if isinstance(vresult, Err):
3016             return vresult
3017
3018         # If the string group is wrapped inside an Atom node, we must make sure
3019         # to later replace that Atom with our new (merged) string leaf.
3020         atom_node = LL[string_idx].parent
3021
3022         # We will place BREAK_MARK in between every two substrings that we
3023         # merge. We will then later go through our final result and use the
3024         # various instances of BREAK_MARK we find to add the right values to
3025         # the custom split map.
3026         BREAK_MARK = "@@@@@ BLACK BREAKPOINT MARKER @@@@@"
3027
3028         QUOTE = LL[string_idx].value[-1]
3029
3030         def make_naked(string: str, string_prefix: str) -> str:
3031             """Strip @string (i.e. make it a "naked" string)
3032
3033             Pre-conditions:
3034                 * assert_is_leaf_string(@string)
3035
3036             Returns:
3037                 A string that is identical to @string except that
3038                 @string_prefix has been stripped, the surrounding QUOTE
3039                 characters have been removed, and any remaining QUOTE
3040                 characters have been escaped.
3041             """
3042             assert_is_leaf_string(string)
3043
3044             RE_EVEN_BACKSLASHES = r"(?:(?<!\\)(?:\\\\)*)"
3045             naked_string = string[len(string_prefix) + 1 : -1]
3046             naked_string = re.sub(
3047                 "(" + RE_EVEN_BACKSLASHES + ")" + QUOTE, r"\1\\" + QUOTE, naked_string
3048             )
3049             return naked_string
3050
3051         # Holds the CustomSplit objects that will later be added to the custom
3052         # split map.
3053         custom_splits = []
3054
3055         # Temporary storage for the 'has_prefix' part of the CustomSplit objects.
3056         prefix_tracker = []
3057
3058         # Sets the 'prefix' variable. This is the prefix that the final merged
3059         # string will have.
3060         next_str_idx = string_idx
3061         prefix = ""
3062         while (
3063             not prefix
3064             and is_valid_index(next_str_idx)
3065             and LL[next_str_idx].type == token.STRING
3066         ):
3067             prefix = get_string_prefix(LL[next_str_idx].value)
3068             next_str_idx += 1
3069
3070         # The next loop merges the string group. The final string will be
3071         # contained in 'S'.
3072         #
3073         # The following convenience variables are used:
3074         #
3075         #   S: string
3076         #   NS: naked string
3077         #   SS: next string
3078         #   NSS: naked next string
3079         S = ""
3080         NS = ""
3081         num_of_strings = 0
3082         next_str_idx = string_idx
3083         while is_valid_index(next_str_idx) and LL[next_str_idx].type == token.STRING:
3084             num_of_strings += 1
3085
3086             SS = LL[next_str_idx].value
3087             next_prefix = get_string_prefix(SS)
3088
3089             # If this is an f-string group but this substring is not prefixed
3090             # with 'f'...
3091             if "f" in prefix and "f" not in next_prefix:
3092                 # Then we must escape any braces contained in this substring.
3093                 SS = re.subf(r"(\{|\})", "{1}{1}", SS)
3094
3095             NSS = make_naked(SS, next_prefix)
3096
3097             has_prefix = bool(next_prefix)
3098             prefix_tracker.append(has_prefix)
3099
3100             S = prefix + QUOTE + NS + NSS + BREAK_MARK + QUOTE
3101             NS = make_naked(S, prefix)
3102
3103             next_str_idx += 1
3104
3105         S_leaf = Leaf(token.STRING, S)
3106         if self.normalize_strings:
3107             normalize_string_quotes(S_leaf)
3108
3109         # Fill the 'custom_splits' list with the appropriate CustomSplit objects.
3110         temp_string = S_leaf.value[len(prefix) + 1 : -1]
3111         for has_prefix in prefix_tracker:
3112             mark_idx = temp_string.find(BREAK_MARK)
3113             assert (
3114                 mark_idx >= 0
3115             ), "Logic error while filling the custom string breakpoint cache."
3116
3117             temp_string = temp_string[mark_idx + len(BREAK_MARK) :]
3118             breakpoint_idx = mark_idx + (len(prefix) if has_prefix else 0) + 1
3119             custom_splits.append(CustomSplit(has_prefix, breakpoint_idx))
3120
3121         string_leaf = Leaf(token.STRING, S_leaf.value.replace(BREAK_MARK, ""))
3122
3123         if atom_node is not None:
3124             replace_child(atom_node, string_leaf)
3125
3126         # Build the final line ('new_line') that this method will later return.
3127         new_line = line.clone()
3128         for (i, leaf) in enumerate(LL):
3129             if i == string_idx:
3130                 new_line.append(string_leaf)
3131
3132             if string_idx <= i < string_idx + num_of_strings:
3133                 for comment_leaf in line.comments_after(LL[i]):
3134                     new_line.append(comment_leaf, preformatted=True)
3135                 continue
3136
3137             append_leaves(new_line, line, [leaf])
3138
3139         self.add_custom_splits(string_leaf.value, custom_splits)
3140         return Ok(new_line)
3141
3142     @staticmethod
3143     def __validate_msg(line: Line, string_idx: int) -> TResult[None]:
3144         """Validate (M)erge (S)tring (G)roup
3145
3146         Transform-time string validation logic for __merge_string_group(...).
3147
3148         Returns:
3149             * Ok(None), if ALL validation checks (listed below) pass.
3150                 OR
3151             * Err(CannotTransform), if any of the following are true:
3152                 - The target string group does not contain ANY stand-alone comments.
3153                 - The target string is not in a string group (i.e. it has no
3154                   adjacent strings).
3155                 - The string group has more than one inline comment.
3156                 - The string group has an inline comment that appears to be a pragma.
3157                 - The set of all string prefixes in the string group is of
3158                   length greater than one and is not equal to {"", "f"}.
3159                 - The string group consists of raw strings.
3160         """
3161         # We first check for "inner" stand-alone comments (i.e. stand-alone
3162         # comments that have a string leaf before them AND after them).
3163         for inc in [1, -1]:
3164             i = string_idx
3165             found_sa_comment = False
3166             is_valid_index = is_valid_index_factory(line.leaves)
3167             while is_valid_index(i) and line.leaves[i].type in [
3168                 token.STRING,
3169                 STANDALONE_COMMENT,
3170             ]:
3171                 if line.leaves[i].type == STANDALONE_COMMENT:
3172                     found_sa_comment = True
3173                 elif found_sa_comment:
3174                     return TErr(
3175                         "StringMerger does NOT merge string groups which contain "
3176                         "stand-alone comments."
3177                     )
3178
3179                 i += inc
3180
3181         num_of_inline_string_comments = 0
3182         set_of_prefixes = set()
3183         num_of_strings = 0
3184         for leaf in line.leaves[string_idx:]:
3185             if leaf.type != token.STRING:
3186                 # If the string group is trailed by a comma, we count the
3187                 # comments trailing the comma to be one of the string group's
3188                 # comments.
3189                 if leaf.type == token.COMMA and id(leaf) in line.comments:
3190                     num_of_inline_string_comments += 1
3191                 break
3192
3193             if has_triple_quotes(leaf.value):
3194                 return TErr("StringMerger does NOT merge multiline strings.")
3195
3196             num_of_strings += 1
3197             prefix = get_string_prefix(leaf.value)
3198             if "r" in prefix:
3199                 return TErr("StringMerger does NOT merge raw strings.")
3200
3201             set_of_prefixes.add(prefix)
3202
3203             if id(leaf) in line.comments:
3204                 num_of_inline_string_comments += 1
3205                 if contains_pragma_comment(line.comments[id(leaf)]):
3206                     return TErr("Cannot merge strings which have pragma comments.")
3207
3208         if num_of_strings < 2:
3209             return TErr(
3210                 f"Not enough strings to merge (num_of_strings={num_of_strings})."
3211             )
3212
3213         if num_of_inline_string_comments > 1:
3214             return TErr(
3215                 f"Too many inline string comments ({num_of_inline_string_comments})."
3216             )
3217
3218         if len(set_of_prefixes) > 1 and set_of_prefixes != {"", "f"}:
3219             return TErr(f"Too many different prefixes ({set_of_prefixes}).")
3220
3221         return Ok(None)
3222
3223
3224 class StringParenStripper(StringTransformer):
3225     """StringTransformer that strips surrounding parentheses from strings.
3226
3227     Requirements:
3228         The line contains a string which is surrounded by parentheses and:
3229             - The target string is NOT the only argument to a function call).
3230             - If the target string contains a PERCENT, the brackets are not
3231               preceeded or followed by an operator with higher precedence than
3232               PERCENT.
3233
3234     Transformations:
3235         The parentheses mentioned in the 'Requirements' section are stripped.
3236
3237     Collaborations:
3238         StringParenStripper has its own inherent usefulness, but it is also
3239         relied on to clean up the parentheses created by StringParenWrapper (in
3240         the event that they are no longer needed).
3241     """
3242
3243     def do_match(self, line: Line) -> TMatchResult:
3244         LL = line.leaves
3245
3246         is_valid_index = is_valid_index_factory(LL)
3247
3248         for (idx, leaf) in enumerate(LL):
3249             # Should be a string...
3250             if leaf.type != token.STRING:
3251                 continue
3252
3253             # Should be preceded by a non-empty LPAR...
3254             if (
3255                 not is_valid_index(idx - 1)
3256                 or LL[idx - 1].type != token.LPAR
3257                 or is_empty_lpar(LL[idx - 1])
3258             ):
3259                 continue
3260
3261             # That LPAR should NOT be preceded by a function name or a closing
3262             # bracket (which could be a function which returns a function or a
3263             # list/dictionary that contains a function)...
3264             if is_valid_index(idx - 2) and (
3265                 LL[idx - 2].type == token.NAME or LL[idx - 2].type in CLOSING_BRACKETS
3266             ):
3267                 continue
3268
3269             string_idx = idx
3270
3271             # Skip the string trailer, if one exists.
3272             string_parser = StringParser()
3273             next_idx = string_parser.parse(LL, string_idx)
3274
3275             # if the leaves in the parsed string include a PERCENT, we need to
3276             # make sure the initial LPAR is NOT preceded by an operator with
3277             # higher or equal precedence to PERCENT
3278             if is_valid_index(idx - 2):
3279                 # mypy can't quite follow unless we name this
3280                 before_lpar = LL[idx - 2]
3281                 if token.PERCENT in {leaf.type for leaf in LL[idx - 1 : next_idx]} and (
3282                     (
3283                         before_lpar.type
3284                         in {
3285                             token.STAR,
3286                             token.AT,
3287                             token.SLASH,
3288                             token.DOUBLESLASH,
3289                             token.PERCENT,
3290                             token.TILDE,
3291                             token.DOUBLESTAR,
3292                             token.AWAIT,
3293                             token.LSQB,
3294                             token.LPAR,
3295                         }
3296                     )
3297                     or (
3298                         # only unary PLUS/MINUS
3299                         before_lpar.parent
3300                         and before_lpar.parent.type == syms.factor
3301                         and (before_lpar.type in {token.PLUS, token.MINUS})
3302                     )
3303                 ):
3304                     continue
3305
3306             # Should be followed by a non-empty RPAR...
3307             if (
3308                 is_valid_index(next_idx)
3309                 and LL[next_idx].type == token.RPAR
3310                 and not is_empty_rpar(LL[next_idx])
3311             ):
3312                 # That RPAR should NOT be followed by anything with higher
3313                 # precedence than PERCENT
3314                 if is_valid_index(next_idx + 1) and LL[next_idx + 1].type in {
3315                     token.DOUBLESTAR,
3316                     token.LSQB,
3317                     token.LPAR,
3318                     token.DOT,
3319                 }:
3320                     continue
3321
3322                 return Ok(string_idx)
3323
3324         return TErr("This line has no strings wrapped in parens.")
3325
3326     def do_transform(self, line: Line, string_idx: int) -> Iterator[TResult[Line]]:
3327         LL = line.leaves
3328
3329         string_parser = StringParser()
3330         rpar_idx = string_parser.parse(LL, string_idx)
3331
3332         for leaf in (LL[string_idx - 1], LL[rpar_idx]):
3333             if line.comments_after(leaf):
3334                 yield TErr(
3335                     "Will not strip parentheses which have comments attached to them."
3336                 )
3337                 return
3338
3339         new_line = line.clone()
3340         new_line.comments = line.comments.copy()
3341         try:
3342             append_leaves(new_line, line, LL[: string_idx - 1])
3343         except BracketMatchError:
3344             # HACK: I believe there is currently a bug somewhere in
3345             # right_hand_split() that is causing brackets to not be tracked
3346             # properly by a shared BracketTracker.
3347             append_leaves(new_line, line, LL[: string_idx - 1], preformatted=True)
3348
3349         string_leaf = Leaf(token.STRING, LL[string_idx].value)
3350         LL[string_idx - 1].remove()
3351         replace_child(LL[string_idx], string_leaf)
3352         new_line.append(string_leaf)
3353
3354         append_leaves(
3355             new_line, line, LL[string_idx + 1 : rpar_idx] + LL[rpar_idx + 1 :]
3356         )
3357
3358         LL[rpar_idx].remove()
3359
3360         yield Ok(new_line)
3361
3362
3363 class BaseStringSplitter(StringTransformer):
3364     """
3365     Abstract class for StringTransformers which transform a Line's strings by splitting
3366     them or placing them on their own lines where necessary to avoid going over
3367     the configured line length.
3368
3369     Requirements:
3370         * The target string value is responsible for the line going over the
3371         line length limit. It follows that after all of black's other line
3372         split methods have been exhausted, this line (or one of the resulting
3373         lines after all line splits are performed) would still be over the
3374         line_length limit unless we split this string.
3375             AND
3376         * The target string is NOT a "pointless" string (i.e. a string that has
3377         no parent or siblings).
3378             AND
3379         * The target string is not followed by an inline comment that appears
3380         to be a pragma.
3381             AND
3382         * The target string is not a multiline (i.e. triple-quote) string.
3383     """
3384
3385     @abstractmethod
3386     def do_splitter_match(self, line: Line) -> TMatchResult:
3387         """
3388         BaseStringSplitter asks its clients to override this method instead of
3389         `StringTransformer.do_match(...)`.
3390
3391         Follows the same protocol as `StringTransformer.do_match(...)`.
3392
3393         Refer to `help(StringTransformer.do_match)` for more information.
3394         """
3395
3396     def do_match(self, line: Line) -> TMatchResult:
3397         match_result = self.do_splitter_match(line)
3398         if isinstance(match_result, Err):
3399             return match_result
3400
3401         string_idx = match_result.ok()
3402         vresult = self.__validate(line, string_idx)
3403         if isinstance(vresult, Err):
3404             return vresult
3405
3406         return match_result
3407
3408     def __validate(self, line: Line, string_idx: int) -> TResult[None]:
3409         """
3410         Checks that @line meets all of the requirements listed in this classes'
3411         docstring. Refer to `help(BaseStringSplitter)` for a detailed
3412         description of those requirements.
3413
3414         Returns:
3415             * Ok(None), if ALL of the requirements are met.
3416                 OR
3417             * Err(CannotTransform), if ANY of the requirements are NOT met.
3418         """
3419         LL = line.leaves
3420
3421         string_leaf = LL[string_idx]
3422
3423         max_string_length = self.__get_max_string_length(line, string_idx)
3424         if len(string_leaf.value) <= max_string_length:
3425             return TErr(
3426                 "The string itself is not what is causing this line to be too long."
3427             )
3428
3429         if not string_leaf.parent or [L.type for L in string_leaf.parent.children] == [
3430             token.STRING,
3431             token.NEWLINE,
3432         ]:
3433             return TErr(
3434                 f"This string ({string_leaf.value}) appears to be pointless (i.e. has"
3435                 " no parent)."
3436             )
3437
3438         if id(line.leaves[string_idx]) in line.comments and contains_pragma_comment(
3439             line.comments[id(line.leaves[string_idx])]
3440         ):
3441             return TErr(
3442                 "Line appears to end with an inline pragma comment. Splitting the line"
3443                 " could modify the pragma's behavior."
3444             )
3445
3446         if has_triple_quotes(string_leaf.value):
3447             return TErr("We cannot split multiline strings.")
3448
3449         return Ok(None)
3450
3451     def __get_max_string_length(self, line: Line, string_idx: int) -> int:
3452         """
3453         Calculates the max string length used when attempting to determine
3454         whether or not the target string is responsible for causing the line to
3455         go over the line length limit.
3456
3457         WARNING: This method is tightly coupled to both StringSplitter and
3458         (especially) StringParenWrapper. There is probably a better way to
3459         accomplish what is being done here.
3460
3461         Returns:
3462             max_string_length: such that `line.leaves[string_idx].value >
3463             max_string_length` implies that the target string IS responsible
3464             for causing this line to exceed the line length limit.
3465         """
3466         LL = line.leaves
3467
3468         is_valid_index = is_valid_index_factory(LL)
3469
3470         # We use the shorthand "WMA4" in comments to abbreviate "We must
3471         # account for". When giving examples, we use STRING to mean some/any
3472         # valid string.
3473         #
3474         # Finally, we use the following convenience variables:
3475         #
3476         #   P:  The leaf that is before the target string leaf.
3477         #   N:  The leaf that is after the target string leaf.
3478         #   NN: The leaf that is after N.
3479
3480         # WMA4 the whitespace at the beginning of the line.
3481         offset = line.depth * 4
3482
3483         if is_valid_index(string_idx - 1):
3484             p_idx = string_idx - 1
3485             if (
3486                 LL[string_idx - 1].type == token.LPAR
3487                 and LL[string_idx - 1].value == ""
3488                 and string_idx >= 2
3489             ):
3490                 # If the previous leaf is an empty LPAR placeholder, we should skip it.
3491                 p_idx -= 1
3492
3493             P = LL[p_idx]
3494             if P.type == token.PLUS:
3495                 # WMA4 a space and a '+' character (e.g. `+ STRING`).
3496                 offset += 2
3497
3498             if P.type == token.COMMA:
3499                 # WMA4 a space, a comma, and a closing bracket [e.g. `), STRING`].
3500                 offset += 3
3501
3502             if P.type in [token.COLON, token.EQUAL, token.NAME]:
3503                 # This conditional branch is meant to handle dictionary keys,
3504                 # variable assignments, 'return STRING' statement lines, and
3505                 # 'else STRING' ternary expression lines.
3506
3507                 # WMA4 a single space.
3508                 offset += 1
3509
3510                 # WMA4 the lengths of any leaves that came before that space,
3511                 # but after any closing bracket before that space.
3512                 for leaf in reversed(LL[: p_idx + 1]):
3513                     offset += len(str(leaf))
3514                     if leaf.type in CLOSING_BRACKETS:
3515                         break
3516
3517         if is_valid_index(string_idx + 1):
3518             N = LL[string_idx + 1]
3519             if N.type == token.RPAR and N.value == "" and len(LL) > string_idx + 2:
3520                 # If the next leaf is an empty RPAR placeholder, we should skip it.
3521                 N = LL[string_idx + 2]
3522
3523             if N.type == token.COMMA:
3524                 # WMA4 a single comma at the end of the string (e.g `STRING,`).
3525                 offset += 1
3526
3527             if is_valid_index(string_idx + 2):
3528                 NN = LL[string_idx + 2]
3529
3530                 if N.type == token.DOT and NN.type == token.NAME:
3531                     # This conditional branch is meant to handle method calls invoked
3532                     # off of a string literal up to and including the LPAR character.
3533
3534                     # WMA4 the '.' character.
3535                     offset += 1
3536
3537                     if (
3538                         is_valid_index(string_idx + 3)
3539                         and LL[string_idx + 3].type == token.LPAR
3540                     ):
3541                         # WMA4 the left parenthesis character.
3542                         offset += 1
3543
3544                     # WMA4 the length of the method's name.
3545                     offset += len(NN.value)
3546
3547         has_comments = False
3548         for comment_leaf in line.comments_after(LL[string_idx]):
3549             if not has_comments:
3550                 has_comments = True
3551                 # WMA4 two spaces before the '#' character.
3552                 offset += 2
3553
3554             # WMA4 the length of the inline comment.
3555             offset += len(comment_leaf.value)
3556
3557         max_string_length = self.line_length - offset
3558         return max_string_length
3559
3560
3561 class StringSplitter(CustomSplitMapMixin, BaseStringSplitter):
3562     """
3563     StringTransformer that splits "atom" strings (i.e. strings which exist on
3564     lines by themselves).
3565
3566     Requirements:
3567         * The line consists ONLY of a single string (with the exception of a
3568         '+' symbol which MAY exist at the start of the line), MAYBE a string
3569         trailer, and MAYBE a trailing comma.
3570             AND
3571         * All of the requirements listed in BaseStringSplitter's docstring.
3572
3573     Transformations:
3574         The string mentioned in the 'Requirements' section is split into as
3575         many substrings as necessary to adhere to the configured line length.
3576
3577         In the final set of substrings, no substring should be smaller than
3578         MIN_SUBSTR_SIZE characters.
3579
3580         The string will ONLY be split on spaces (i.e. each new substring should
3581         start with a space).
3582
3583         If the string is an f-string, it will NOT be split in the middle of an
3584         f-expression (e.g. in f"FooBar: {foo() if x else bar()}", {foo() if x
3585         else bar()} is an f-expression).
3586
3587         If the string that is being split has an associated set of custom split
3588         records and those custom splits will NOT result in any line going over
3589         the configured line length, those custom splits are used. Otherwise the
3590         string is split as late as possible (from left-to-right) while still
3591         adhering to the transformation rules listed above.
3592
3593     Collaborations:
3594         StringSplitter relies on StringMerger to construct the appropriate
3595         CustomSplit objects and add them to the custom split map.
3596     """
3597
3598     MIN_SUBSTR_SIZE = 6
3599     # Matches an "f-expression" (e.g. {var}) that might be found in an f-string.
3600     RE_FEXPR = r"""
3601     (?<!\{)\{
3602         (?:
3603             [^\{\}]
3604             | \{\{
3605             | \}\}
3606         )+?
3607     (?<!\})(?:\}\})*\}(?!\})
3608     """
3609
3610     def do_splitter_match(self, line: Line) -> TMatchResult:
3611         LL = line.leaves
3612
3613         is_valid_index = is_valid_index_factory(LL)
3614
3615         idx = 0
3616
3617         # The first leaf MAY be a '+' symbol...
3618         if is_valid_index(idx) and LL[idx].type == token.PLUS:
3619             idx += 1
3620
3621         # The next/first leaf MAY be an empty LPAR...
3622         if is_valid_index(idx) and is_empty_lpar(LL[idx]):
3623             idx += 1
3624
3625         # The next/first leaf MUST be a string...
3626         if not is_valid_index(idx) or LL[idx].type != token.STRING:
3627             return TErr("Line does not start with a string.")
3628
3629         string_idx = idx
3630
3631         # Skip the string trailer, if one exists.
3632         string_parser = StringParser()
3633         idx = string_parser.parse(LL, string_idx)
3634
3635         # That string MAY be followed by an empty RPAR...
3636         if is_valid_index(idx) and is_empty_rpar(LL[idx]):
3637             idx += 1
3638
3639         # That string / empty RPAR leaf MAY be followed by a comma...
3640         if is_valid_index(idx) and LL[idx].type == token.COMMA:
3641             idx += 1
3642
3643         # But no more leaves are allowed...
3644         if is_valid_index(idx):
3645             return TErr("This line does not end with a string.")
3646
3647         return Ok(string_idx)
3648
3649     def do_transform(self, line: Line, string_idx: int) -> Iterator[TResult[Line]]:
3650         LL = line.leaves
3651
3652         QUOTE = LL[string_idx].value[-1]
3653
3654         is_valid_index = is_valid_index_factory(LL)
3655         insert_str_child = insert_str_child_factory(LL[string_idx])
3656
3657         prefix = get_string_prefix(LL[string_idx].value)
3658
3659         # We MAY choose to drop the 'f' prefix from substrings that don't
3660         # contain any f-expressions, but ONLY if the original f-string
3661         # contains at least one f-expression. Otherwise, we will alter the AST
3662         # of the program.
3663         drop_pointless_f_prefix = ("f" in prefix) and re.search(
3664             self.RE_FEXPR, LL[string_idx].value, re.VERBOSE
3665         )
3666
3667         first_string_line = True
3668         starts_with_plus = LL[0].type == token.PLUS
3669
3670         def line_needs_plus() -> bool:
3671             return first_string_line and starts_with_plus
3672
3673         def maybe_append_plus(new_line: Line) -> None:
3674             """
3675             Side Effects:
3676                 If @line starts with a plus and this is the first line we are
3677                 constructing, this function appends a PLUS leaf to @new_line
3678                 and replaces the old PLUS leaf in the node structure. Otherwise
3679                 this function does nothing.
3680             """
3681             if line_needs_plus():
3682                 plus_leaf = Leaf(token.PLUS, "+")
3683                 replace_child(LL[0], plus_leaf)
3684                 new_line.append(plus_leaf)
3685
3686         ends_with_comma = (
3687             is_valid_index(string_idx + 1) and LL[string_idx + 1].type == token.COMMA
3688         )
3689
3690         def max_last_string() -> int:
3691             """
3692             Returns:
3693                 The max allowed length of the string value used for the last
3694                 line we will construct.
3695             """
3696             result = self.line_length
3697             result -= line.depth * 4
3698             result -= 1 if ends_with_comma else 0
3699             result -= 2 if line_needs_plus() else 0
3700             return result
3701
3702         # --- Calculate Max Break Index (for string value)
3703         # We start with the line length limit
3704         max_break_idx = self.line_length
3705         # The last index of a string of length N is N-1.
3706         max_break_idx -= 1
3707         # Leading whitespace is not present in the string value (e.g. Leaf.value).
3708         max_break_idx -= line.depth * 4
3709         if max_break_idx < 0:
3710             yield TErr(
3711                 f"Unable to split {LL[string_idx].value} at such high of a line depth:"
3712                 f" {line.depth}"
3713             )
3714             return
3715
3716         # Check if StringMerger registered any custom splits.
3717         custom_splits = self.pop_custom_splits(LL[string_idx].value)
3718         # We use them ONLY if none of them would produce lines that exceed the
3719         # line limit.
3720         use_custom_breakpoints = bool(
3721             custom_splits
3722             and all(csplit.break_idx <= max_break_idx for csplit in custom_splits)
3723         )
3724
3725         # Temporary storage for the remaining chunk of the string line that
3726         # can't fit onto the line currently being constructed.
3727         rest_value = LL[string_idx].value
3728
3729         def more_splits_should_be_made() -> bool:
3730             """
3731             Returns:
3732                 True iff `rest_value` (the remaining string value from the last
3733                 split), should be split again.
3734             """
3735             if use_custom_breakpoints:
3736                 return len(custom_splits) > 1
3737             else:
3738                 return len(rest_value) > max_last_string()
3739
3740         string_line_results: List[Ok[Line]] = []
3741         while more_splits_should_be_made():
3742             if use_custom_breakpoints:
3743                 # Custom User Split (manual)
3744                 csplit = custom_splits.pop(0)
3745                 break_idx = csplit.break_idx
3746             else:
3747                 # Algorithmic Split (automatic)
3748                 max_bidx = max_break_idx - 2 if line_needs_plus() else max_break_idx
3749                 maybe_break_idx = self.__get_break_idx(rest_value, max_bidx)
3750                 if maybe_break_idx is None:
3751                     # If we are unable to algorithmically determine a good split
3752                     # and this string has custom splits registered to it, we
3753                     # fall back to using them--which means we have to start
3754                     # over from the beginning.
3755                     if custom_splits:
3756                         rest_value = LL[string_idx].value
3757                         string_line_results = []
3758                         first_string_line = True
3759                         use_custom_breakpoints = True
3760                         continue
3761
3762                     # Otherwise, we stop splitting here.
3763                     break
3764
3765                 break_idx = maybe_break_idx
3766
3767             # --- Construct `next_value`
3768             next_value = rest_value[:break_idx] + QUOTE
3769             if (
3770                 # Are we allowed to try to drop a pointless 'f' prefix?
3771                 drop_pointless_f_prefix
3772                 # If we are, will we be successful?
3773                 and next_value != self.__normalize_f_string(next_value, prefix)
3774             ):
3775                 # If the current custom split did NOT originally use a prefix,
3776                 # then `csplit.break_idx` will be off by one after removing
3777                 # the 'f' prefix.
3778                 break_idx = (
3779                     break_idx + 1
3780                     if use_custom_breakpoints and not csplit.has_prefix
3781                     else break_idx
3782                 )
3783                 next_value = rest_value[:break_idx] + QUOTE
3784                 next_value = self.__normalize_f_string(next_value, prefix)
3785
3786             # --- Construct `next_leaf`
3787             next_leaf = Leaf(token.STRING, next_value)
3788             insert_str_child(next_leaf)
3789             self.__maybe_normalize_string_quotes(next_leaf)
3790
3791             # --- Construct `next_line`
3792             next_line = line.clone()
3793             maybe_append_plus(next_line)
3794             next_line.append(next_leaf)
3795             string_line_results.append(Ok(next_line))
3796
3797             rest_value = prefix + QUOTE + rest_value[break_idx:]
3798             first_string_line = False
3799
3800         yield from string_line_results
3801
3802         if drop_pointless_f_prefix:
3803             rest_value = self.__normalize_f_string(rest_value, prefix)
3804
3805         rest_leaf = Leaf(token.STRING, rest_value)
3806         insert_str_child(rest_leaf)
3807
3808         # NOTE: I could not find a test case that verifies that the following
3809         # line is actually necessary, but it seems to be. Otherwise we risk
3810         # not normalizing the last substring, right?
3811         self.__maybe_normalize_string_quotes(rest_leaf)
3812
3813         last_line = line.clone()
3814         maybe_append_plus(last_line)
3815
3816         # If there are any leaves to the right of the target string...
3817         if is_valid_index(string_idx + 1):
3818             # We use `temp_value` here to determine how long the last line
3819             # would be if we were to append all the leaves to the right of the
3820             # target string to the last string line.
3821             temp_value = rest_value
3822             for leaf in LL[string_idx + 1 :]:
3823                 temp_value += str(leaf)
3824                 if leaf.type == token.LPAR:
3825                     break
3826
3827             # Try to fit them all on the same line with the last substring...
3828             if (
3829                 len(temp_value) <= max_last_string()
3830                 or LL[string_idx + 1].type == token.COMMA
3831             ):
3832                 last_line.append(rest_leaf)
3833                 append_leaves(last_line, line, LL[string_idx + 1 :])
3834                 yield Ok(last_line)
3835             # Otherwise, place the last substring on one line and everything
3836             # else on a line below that...
3837             else:
3838                 last_line.append(rest_leaf)
3839                 yield Ok(last_line)
3840
3841                 non_string_line = line.clone()
3842                 append_leaves(non_string_line, line, LL[string_idx + 1 :])
3843                 yield Ok(non_string_line)
3844         # Else the target string was the last leaf...
3845         else:
3846             last_line.append(rest_leaf)
3847             last_line.comments = line.comments.copy()
3848             yield Ok(last_line)
3849
3850     def __get_break_idx(self, string: str, max_break_idx: int) -> Optional[int]:
3851         """
3852         This method contains the algorithm that StringSplitter uses to
3853         determine which character to split each string at.
3854
3855         Args:
3856             @string: The substring that we are attempting to split.
3857             @max_break_idx: The ideal break index. We will return this value if it
3858             meets all the necessary conditions. In the likely event that it
3859             doesn't we will try to find the closest index BELOW @max_break_idx
3860             that does. If that fails, we will expand our search by also
3861             considering all valid indices ABOVE @max_break_idx.
3862
3863         Pre-Conditions:
3864             * assert_is_leaf_string(@string)
3865             * 0 <= @max_break_idx < len(@string)
3866
3867         Returns:
3868             break_idx, if an index is able to be found that meets all of the
3869             conditions listed in the 'Transformations' section of this classes'
3870             docstring.
3871                 OR
3872             None, otherwise.
3873         """
3874         is_valid_index = is_valid_index_factory(string)
3875
3876         assert is_valid_index(max_break_idx)
3877         assert_is_leaf_string(string)
3878
3879         _fexpr_slices: Optional[List[Tuple[Index, Index]]] = None
3880
3881         def fexpr_slices() -> Iterator[Tuple[Index, Index]]:
3882             """
3883             Yields:
3884                 All ranges of @string which, if @string were to be split there,
3885                 would result in the splitting of an f-expression (which is NOT
3886                 allowed).
3887             """
3888             nonlocal _fexpr_slices
3889
3890             if _fexpr_slices is None:
3891                 _fexpr_slices = []
3892                 for match in re.finditer(self.RE_FEXPR, string, re.VERBOSE):
3893                     _fexpr_slices.append(match.span())
3894
3895             yield from _fexpr_slices
3896
3897         is_fstring = "f" in get_string_prefix(string)
3898
3899         def breaks_fstring_expression(i: Index) -> bool:
3900             """
3901             Returns:
3902                 True iff returning @i would result in the splitting of an
3903                 f-expression (which is NOT allowed).
3904             """
3905             if not is_fstring:
3906                 return False
3907
3908             for (start, end) in fexpr_slices():
3909                 if start <= i < end:
3910                     return True
3911
3912             return False
3913
3914         def passes_all_checks(i: Index) -> bool:
3915             """
3916             Returns:
3917                 True iff ALL of the conditions listed in the 'Transformations'
3918                 section of this classes' docstring would be be met by returning @i.
3919             """
3920             is_space = string[i] == " "
3921             is_big_enough = (
3922                 len(string[i:]) >= self.MIN_SUBSTR_SIZE
3923                 and len(string[:i]) >= self.MIN_SUBSTR_SIZE
3924             )
3925             return is_space and is_big_enough and not breaks_fstring_expression(i)
3926
3927         # First, we check all indices BELOW @max_break_idx.
3928         break_idx = max_break_idx
3929         while is_valid_index(break_idx - 1) and not passes_all_checks(break_idx):
3930             break_idx -= 1
3931
3932         if not passes_all_checks(break_idx):
3933             # If that fails, we check all indices ABOVE @max_break_idx.
3934             #
3935             # If we are able to find a valid index here, the next line is going
3936             # to be longer than the specified line length, but it's probably
3937             # better than doing nothing at all.
3938             break_idx = max_break_idx + 1
3939             while is_valid_index(break_idx + 1) and not passes_all_checks(break_idx):
3940                 break_idx += 1
3941
3942             if not is_valid_index(break_idx) or not passes_all_checks(break_idx):
3943                 return None
3944
3945         return break_idx
3946
3947     def __maybe_normalize_string_quotes(self, leaf: Leaf) -> None:
3948         if self.normalize_strings:
3949             normalize_string_quotes(leaf)
3950
3951     def __normalize_f_string(self, string: str, prefix: str) -> str:
3952         """
3953         Pre-Conditions:
3954             * assert_is_leaf_string(@string)
3955
3956         Returns:
3957             * If @string is an f-string that contains no f-expressions, we
3958             return a string identical to @string except that the 'f' prefix
3959             has been stripped and all double braces (i.e. '{{' or '}}') have
3960             been normalized (i.e. turned into '{' or '}').
3961                 OR
3962             * Otherwise, we return @string.
3963         """
3964         assert_is_leaf_string(string)
3965
3966         if "f" in prefix and not re.search(self.RE_FEXPR, string, re.VERBOSE):
3967             new_prefix = prefix.replace("f", "")
3968
3969             temp = string[len(prefix) :]
3970             temp = re.sub(r"\{\{", "{", temp)
3971             temp = re.sub(r"\}\}", "}", temp)
3972             new_string = temp
3973
3974             return f"{new_prefix}{new_string}"
3975         else:
3976             return string
3977
3978
3979 class StringParenWrapper(CustomSplitMapMixin, BaseStringSplitter):
3980     """
3981     StringTransformer that splits non-"atom" strings (i.e. strings that do not
3982     exist on lines by themselves).
3983
3984     Requirements:
3985         All of the requirements listed in BaseStringSplitter's docstring in
3986         addition to the requirements listed below:
3987
3988         * The line is a return/yield statement, which returns/yields a string.
3989             OR
3990         * The line is part of a ternary expression (e.g. `x = y if cond else
3991         z`) such that the line starts with `else <string>`, where <string> is
3992         some string.
3993             OR
3994         * The line is an assert statement, which ends with a string.
3995             OR
3996         * The line is an assignment statement (e.g. `x = <string>` or `x +=
3997         <string>`) such that the variable is being assigned the value of some
3998         string.
3999             OR
4000         * The line is a dictionary key assignment where some valid key is being
4001         assigned the value of some string.
4002
4003     Transformations:
4004         The chosen string is wrapped in parentheses and then split at the LPAR.
4005
4006         We then have one line which ends with an LPAR and another line that
4007         starts with the chosen string. The latter line is then split again at
4008         the RPAR. This results in the RPAR (and possibly a trailing comma)
4009         being placed on its own line.
4010
4011         NOTE: If any leaves exist to the right of the chosen string (except
4012         for a trailing comma, which would be placed after the RPAR), those
4013         leaves are placed inside the parentheses.  In effect, the chosen
4014         string is not necessarily being "wrapped" by parentheses. We can,
4015         however, count on the LPAR being placed directly before the chosen
4016         string.
4017
4018         In other words, StringParenWrapper creates "atom" strings. These
4019         can then be split again by StringSplitter, if necessary.
4020
4021     Collaborations:
4022         In the event that a string line split by StringParenWrapper is
4023         changed such that it no longer needs to be given its own line,
4024         StringParenWrapper relies on StringParenStripper to clean up the
4025         parentheses it created.
4026     """
4027
4028     def do_splitter_match(self, line: Line) -> TMatchResult:
4029         LL = line.leaves
4030
4031         string_idx = (
4032             self._return_match(LL)
4033             or self._else_match(LL)
4034             or self._assert_match(LL)
4035             or self._assign_match(LL)
4036             or self._dict_match(LL)
4037         )
4038
4039         if string_idx is not None:
4040             string_value = line.leaves[string_idx].value
4041             # If the string has no spaces...
4042             if " " not in string_value:
4043                 # And will still violate the line length limit when split...
4044                 max_string_length = self.line_length - ((line.depth + 1) * 4)
4045                 if len(string_value) > max_string_length:
4046                     # And has no associated custom splits...
4047                     if not self.has_custom_splits(string_value):
4048                         # Then we should NOT put this string on its own line.
4049                         return TErr(
4050                             "We do not wrap long strings in parentheses when the"
4051                             " resultant line would still be over the specified line"
4052                             " length and can't be split further by StringSplitter."
4053                         )
4054             return Ok(string_idx)
4055
4056         return TErr("This line does not contain any non-atomic strings.")
4057
4058     @staticmethod
4059     def _return_match(LL: List[Leaf]) -> Optional[int]:
4060         """
4061         Returns:
4062             string_idx such that @LL[string_idx] is equal to our target (i.e.
4063             matched) string, if this line matches the return/yield statement
4064             requirements listed in the 'Requirements' section of this classes'
4065             docstring.
4066                 OR
4067             None, otherwise.
4068         """
4069         # If this line is apart of a return/yield statement and the first leaf
4070         # contains either the "return" or "yield" keywords...
4071         if parent_type(LL[0]) in [syms.return_stmt, syms.yield_expr] and LL[
4072             0
4073         ].value in ["return", "yield"]:
4074             is_valid_index = is_valid_index_factory(LL)
4075
4076             idx = 2 if is_valid_index(1) and is_empty_par(LL[1]) else 1
4077             # The next visible leaf MUST contain a string...
4078             if is_valid_index(idx) and LL[idx].type == token.STRING:
4079                 return idx
4080
4081         return None
4082
4083     @staticmethod
4084     def _else_match(LL: List[Leaf]) -> Optional[int]:
4085         """
4086         Returns:
4087             string_idx such that @LL[string_idx] is equal to our target (i.e.
4088             matched) string, if this line matches the ternary expression
4089             requirements listed in the 'Requirements' section of this classes'
4090             docstring.
4091                 OR
4092             None, otherwise.
4093         """
4094         # If this line is apart of a ternary expression and the first leaf
4095         # contains the "else" keyword...
4096         if (
4097             parent_type(LL[0]) == syms.test
4098             and LL[0].type == token.NAME
4099             and LL[0].value == "else"
4100         ):
4101             is_valid_index = is_valid_index_factory(LL)
4102
4103             idx = 2 if is_valid_index(1) and is_empty_par(LL[1]) else 1
4104             # The next visible leaf MUST contain a string...
4105             if is_valid_index(idx) and LL[idx].type == token.STRING:
4106                 return idx
4107
4108         return None
4109
4110     @staticmethod
4111     def _assert_match(LL: List[Leaf]) -> Optional[int]:
4112         """
4113         Returns:
4114             string_idx such that @LL[string_idx] is equal to our target (i.e.
4115             matched) string, if this line matches the assert statement
4116             requirements listed in the 'Requirements' section of this classes'
4117             docstring.
4118                 OR
4119             None, otherwise.
4120         """
4121         # If this line is apart of an assert statement and the first leaf
4122         # contains the "assert" keyword...
4123         if parent_type(LL[0]) == syms.assert_stmt and LL[0].value == "assert":
4124             is_valid_index = is_valid_index_factory(LL)
4125
4126             for (i, leaf) in enumerate(LL):
4127                 # We MUST find a comma...
4128                 if leaf.type == token.COMMA:
4129                     idx = i + 2 if is_empty_par(LL[i + 1]) else i + 1
4130
4131                     # That comma MUST be followed by a string...
4132                     if is_valid_index(idx) and LL[idx].type == token.STRING:
4133                         string_idx = idx
4134
4135                         # Skip the string trailer, if one exists.
4136                         string_parser = StringParser()
4137                         idx = string_parser.parse(LL, string_idx)
4138
4139                         # But no more leaves are allowed...
4140                         if not is_valid_index(idx):
4141                             return string_idx
4142
4143         return None
4144
4145     @staticmethod
4146     def _assign_match(LL: List[Leaf]) -> Optional[int]:
4147         """
4148         Returns:
4149             string_idx such that @LL[string_idx] is equal to our target (i.e.
4150             matched) string, if this line matches the assignment statement
4151             requirements listed in the 'Requirements' section of this classes'
4152             docstring.
4153                 OR
4154             None, otherwise.
4155         """
4156         # If this line is apart of an expression statement or is a function
4157         # argument AND the first leaf contains a variable name...
4158         if (
4159             parent_type(LL[0]) in [syms.expr_stmt, syms.argument, syms.power]
4160             and LL[0].type == token.NAME
4161         ):
4162             is_valid_index = is_valid_index_factory(LL)
4163
4164             for (i, leaf) in enumerate(LL):
4165                 # We MUST find either an '=' or '+=' symbol...
4166                 if leaf.type in [token.EQUAL, token.PLUSEQUAL]:
4167                     idx = i + 2 if is_empty_par(LL[i + 1]) else i + 1
4168
4169                     # That symbol MUST be followed by a string...
4170                     if is_valid_index(idx) and LL[idx].type == token.STRING:
4171                         string_idx = idx
4172
4173                         # Skip the string trailer, if one exists.
4174                         string_parser = StringParser()
4175                         idx = string_parser.parse(LL, string_idx)
4176
4177                         # The next leaf MAY be a comma iff this line is apart
4178                         # of a function argument...
4179                         if (
4180                             parent_type(LL[0]) == syms.argument
4181                             and is_valid_index(idx)
4182                             and LL[idx].type == token.COMMA
4183                         ):
4184                             idx += 1
4185
4186                         # But no more leaves are allowed...
4187                         if not is_valid_index(idx):
4188                             return string_idx
4189
4190         return None
4191
4192     @staticmethod
4193     def _dict_match(LL: List[Leaf]) -> Optional[int]:
4194         """
4195         Returns:
4196             string_idx such that @LL[string_idx] is equal to our target (i.e.
4197             matched) string, if this line matches the dictionary key assignment
4198             statement requirements listed in the 'Requirements' section of this
4199             classes' docstring.
4200                 OR
4201             None, otherwise.
4202         """
4203         # If this line is apart of a dictionary key assignment...
4204         if syms.dictsetmaker in [parent_type(LL[0]), parent_type(LL[0].parent)]:
4205             is_valid_index = is_valid_index_factory(LL)
4206
4207             for (i, leaf) in enumerate(LL):
4208                 # We MUST find a colon...
4209                 if leaf.type == token.COLON:
4210                     idx = i + 2 if is_empty_par(LL[i + 1]) else i + 1
4211
4212                     # That colon MUST be followed by a string...
4213                     if is_valid_index(idx) and LL[idx].type == token.STRING:
4214                         string_idx = idx
4215
4216                         # Skip the string trailer, if one exists.
4217                         string_parser = StringParser()
4218                         idx = string_parser.parse(LL, string_idx)
4219
4220                         # That string MAY be followed by a comma...
4221                         if is_valid_index(idx) and LL[idx].type == token.COMMA:
4222                             idx += 1
4223
4224                         # But no more leaves are allowed...
4225                         if not is_valid_index(idx):
4226                             return string_idx
4227
4228         return None
4229
4230     def do_transform(self, line: Line, string_idx: int) -> Iterator[TResult[Line]]:
4231         LL = line.leaves
4232
4233         is_valid_index = is_valid_index_factory(LL)
4234         insert_str_child = insert_str_child_factory(LL[string_idx])
4235
4236         comma_idx = -1
4237         ends_with_comma = False
4238         if LL[comma_idx].type == token.COMMA:
4239             ends_with_comma = True
4240
4241         leaves_to_steal_comments_from = [LL[string_idx]]
4242         if ends_with_comma:
4243             leaves_to_steal_comments_from.append(LL[comma_idx])
4244
4245         # --- First Line
4246         first_line = line.clone()
4247         left_leaves = LL[:string_idx]
4248
4249         # We have to remember to account for (possibly invisible) LPAR and RPAR
4250         # leaves that already wrapped the target string. If these leaves do
4251         # exist, we will replace them with our own LPAR and RPAR leaves.
4252         old_parens_exist = False
4253         if left_leaves and left_leaves[-1].type == token.LPAR:
4254             old_parens_exist = True
4255             leaves_to_steal_comments_from.append(left_leaves[-1])
4256             left_leaves.pop()
4257
4258         append_leaves(first_line, line, left_leaves)
4259
4260         lpar_leaf = Leaf(token.LPAR, "(")
4261         if old_parens_exist:
4262             replace_child(LL[string_idx - 1], lpar_leaf)
4263         else:
4264             insert_str_child(lpar_leaf)
4265         first_line.append(lpar_leaf)
4266
4267         # We throw inline comments that were originally to the right of the
4268         # target string to the top line. They will now be shown to the right of
4269         # the LPAR.
4270         for leaf in leaves_to_steal_comments_from:
4271             for comment_leaf in line.comments_after(leaf):
4272                 first_line.append(comment_leaf, preformatted=True)
4273
4274         yield Ok(first_line)
4275
4276         # --- Middle (String) Line
4277         # We only need to yield one (possibly too long) string line, since the
4278         # `StringSplitter` will break it down further if necessary.
4279         string_value = LL[string_idx].value
4280         string_line = Line(
4281             depth=line.depth + 1,
4282             inside_brackets=True,
4283             should_explode=line.should_explode,
4284         )
4285         string_leaf = Leaf(token.STRING, string_value)
4286         insert_str_child(string_leaf)
4287         string_line.append(string_leaf)
4288
4289         old_rpar_leaf = None
4290         if is_valid_index(string_idx + 1):
4291             right_leaves = LL[string_idx + 1 :]
4292             if ends_with_comma:
4293                 right_leaves.pop()
4294
4295             if old_parens_exist:
4296                 assert (
4297                     right_leaves and right_leaves[-1].type == token.RPAR
4298                 ), "Apparently, old parentheses do NOT exist?!"
4299                 old_rpar_leaf = right_leaves.pop()
4300
4301             append_leaves(string_line, line, right_leaves)
4302
4303         yield Ok(string_line)
4304
4305         # --- Last Line
4306         last_line = line.clone()
4307         last_line.bracket_tracker = first_line.bracket_tracker
4308
4309         new_rpar_leaf = Leaf(token.RPAR, ")")
4310         if old_rpar_leaf is not None:
4311             replace_child(old_rpar_leaf, new_rpar_leaf)
4312         else:
4313             insert_str_child(new_rpar_leaf)
4314         last_line.append(new_rpar_leaf)
4315
4316         # If the target string ended with a comma, we place this comma to the
4317         # right of the RPAR on the last line.
4318         if ends_with_comma:
4319             comma_leaf = Leaf(token.COMMA, ",")
4320             replace_child(LL[comma_idx], comma_leaf)
4321             last_line.append(comma_leaf)
4322
4323         yield Ok(last_line)
4324
4325
4326 class StringParser:
4327     """
4328     A state machine that aids in parsing a string's "trailer", which can be
4329     either non-existent, an old-style formatting sequence (e.g. `% varX` or `%
4330     (varX, varY)`), or a method-call / attribute access (e.g. `.format(varX,
4331     varY)`).
4332
4333     NOTE: A new StringParser object MUST be instantiated for each string
4334     trailer we need to parse.
4335
4336     Examples:
4337         We shall assume that `line` equals the `Line` object that corresponds
4338         to the following line of python code:
4339         ```
4340         x = "Some {}.".format("String") + some_other_string
4341         ```
4342
4343         Furthermore, we will assume that `string_idx` is some index such that:
4344         ```
4345         assert line.leaves[string_idx].value == "Some {}."
4346         ```
4347
4348         The following code snippet then holds:
4349         ```
4350         string_parser = StringParser()
4351         idx = string_parser.parse(line.leaves, string_idx)
4352         assert line.leaves[idx].type == token.PLUS
4353         ```
4354     """
4355
4356     DEFAULT_TOKEN = -1
4357
4358     # String Parser States
4359     START = 1
4360     DOT = 2
4361     NAME = 3
4362     PERCENT = 4
4363     SINGLE_FMT_ARG = 5
4364     LPAR = 6
4365     RPAR = 7
4366     DONE = 8
4367
4368     # Lookup Table for Next State
4369     _goto: Dict[Tuple[ParserState, NodeType], ParserState] = {
4370         # A string trailer may start with '.' OR '%'.
4371         (START, token.DOT): DOT,
4372         (START, token.PERCENT): PERCENT,
4373         (START, DEFAULT_TOKEN): DONE,
4374         # A '.' MUST be followed by an attribute or method name.
4375         (DOT, token.NAME): NAME,
4376         # A method name MUST be followed by an '(', whereas an attribute name
4377         # is the last symbol in the string trailer.
4378         (NAME, token.LPAR): LPAR,
4379         (NAME, DEFAULT_TOKEN): DONE,
4380         # A '%' symbol can be followed by an '(' or a single argument (e.g. a
4381         # string or variable name).
4382         (PERCENT, token.LPAR): LPAR,
4383         (PERCENT, DEFAULT_TOKEN): SINGLE_FMT_ARG,
4384         # If a '%' symbol is followed by a single argument, that argument is
4385         # the last leaf in the string trailer.
4386         (SINGLE_FMT_ARG, DEFAULT_TOKEN): DONE,
4387         # If present, a ')' symbol is the last symbol in a string trailer.
4388         # (NOTE: LPARS and nested RPARS are not included in this lookup table,
4389         # since they are treated as a special case by the parsing logic in this
4390         # classes' implementation.)
4391         (RPAR, DEFAULT_TOKEN): DONE,
4392     }
4393
4394     def __init__(self) -> None:
4395         self._state = self.START
4396         self._unmatched_lpars = 0
4397
4398     def parse(self, leaves: List[Leaf], string_idx: int) -> int:
4399         """
4400         Pre-conditions:
4401             * @leaves[@string_idx].type == token.STRING
4402
4403         Returns:
4404             The index directly after the last leaf which is apart of the string
4405             trailer, if a "trailer" exists.
4406                 OR
4407             @string_idx + 1, if no string "trailer" exists.
4408         """
4409         assert leaves[string_idx].type == token.STRING
4410
4411         idx = string_idx + 1
4412         while idx < len(leaves) and self._next_state(leaves[idx]):
4413             idx += 1
4414         return idx
4415
4416     def _next_state(self, leaf: Leaf) -> bool:
4417         """
4418         Pre-conditions:
4419             * On the first call to this function, @leaf MUST be the leaf that
4420             was directly after the string leaf in question (e.g. if our target
4421             string is `line.leaves[i]` then the first call to this method must
4422             be `line.leaves[i + 1]`).
4423             * On the next call to this function, the leaf parameter passed in
4424             MUST be the leaf directly following @leaf.
4425
4426         Returns:
4427             True iff @leaf is apart of the string's trailer.
4428         """
4429         # We ignore empty LPAR or RPAR leaves.
4430         if is_empty_par(leaf):
4431             return True
4432
4433         next_token = leaf.type
4434         if next_token == token.LPAR:
4435             self._unmatched_lpars += 1
4436
4437         current_state = self._state
4438
4439         # The LPAR parser state is a special case. We will return True until we
4440         # find the matching RPAR token.
4441         if current_state == self.LPAR:
4442             if next_token == token.RPAR:
4443                 self._unmatched_lpars -= 1
4444                 if self._unmatched_lpars == 0:
4445                     self._state = self.RPAR
4446         # Otherwise, we use a lookup table to determine the next state.
4447         else:
4448             # If the lookup table matches the current state to the next
4449             # token, we use the lookup table.
4450             if (current_state, next_token) in self._goto:
4451                 self._state = self._goto[current_state, next_token]
4452             else:
4453                 # Otherwise, we check if a the current state was assigned a
4454                 # default.
4455                 if (current_state, self.DEFAULT_TOKEN) in self._goto:
4456                     self._state = self._goto[current_state, self.DEFAULT_TOKEN]
4457                 # If no default has been assigned, then this parser has a logic
4458                 # error.
4459                 else:
4460                     raise RuntimeError(f"{self.__class__.__name__} LOGIC ERROR!")
4461
4462             if self._state == self.DONE:
4463                 return False
4464
4465         return True
4466
4467
4468 def TErr(err_msg: str) -> Err[CannotTransform]:
4469     """(T)ransform Err
4470
4471     Convenience function used when working with the TResult type.
4472     """
4473     cant_transform = CannotTransform(err_msg)
4474     return Err(cant_transform)
4475
4476
4477 def contains_pragma_comment(comment_list: List[Leaf]) -> bool:
4478     """
4479     Returns:
4480         True iff one of the comments in @comment_list is a pragma used by one
4481         of the more common static analysis tools for python (e.g. mypy, flake8,
4482         pylint).
4483     """
4484     for comment in comment_list:
4485         if comment.value.startswith(("# type:", "# noqa", "# pylint:")):
4486             return True
4487
4488     return False
4489
4490
4491 def insert_str_child_factory(string_leaf: Leaf) -> Callable[[LN], None]:
4492     """
4493     Factory for a convenience function that is used to orphan @string_leaf
4494     and then insert multiple new leaves into the same part of the node
4495     structure that @string_leaf had originally occupied.
4496
4497     Examples:
4498         Let `string_leaf = Leaf(token.STRING, '"foo"')` and `N =
4499         string_leaf.parent`. Assume the node `N` has the following
4500         original structure:
4501
4502         Node(
4503             expr_stmt, [
4504                 Leaf(NAME, 'x'),
4505                 Leaf(EQUAL, '='),
4506                 Leaf(STRING, '"foo"'),
4507             ]
4508         )
4509
4510         We then run the code snippet shown below.
4511         ```
4512         insert_str_child = insert_str_child_factory(string_leaf)
4513
4514         lpar = Leaf(token.LPAR, '(')
4515         insert_str_child(lpar)
4516
4517         bar = Leaf(token.STRING, '"bar"')
4518         insert_str_child(bar)
4519
4520         rpar = Leaf(token.RPAR, ')')
4521         insert_str_child(rpar)
4522         ```
4523
4524         After which point, it follows that `string_leaf.parent is None` and
4525         the node `N` now has the following structure:
4526
4527         Node(
4528             expr_stmt, [
4529                 Leaf(NAME, 'x'),
4530                 Leaf(EQUAL, '='),
4531                 Leaf(LPAR, '('),
4532                 Leaf(STRING, '"bar"'),
4533                 Leaf(RPAR, ')'),
4534             ]
4535         )
4536     """
4537     string_parent = string_leaf.parent
4538     string_child_idx = string_leaf.remove()
4539
4540     def insert_str_child(child: LN) -> None:
4541         nonlocal string_child_idx
4542
4543         assert string_parent is not None
4544         assert string_child_idx is not None
4545
4546         string_parent.insert_child(string_child_idx, child)
4547         string_child_idx += 1
4548
4549     return insert_str_child
4550
4551
4552 def has_triple_quotes(string: str) -> bool:
4553     """
4554     Returns:
4555         True iff @string starts with three quotation characters.
4556     """
4557     raw_string = string.lstrip(STRING_PREFIX_CHARS)
4558     return raw_string[:3] in {'"""', "'''"}
4559
4560
4561 def parent_type(node: Optional[LN]) -> Optional[NodeType]:
4562     """
4563     Returns:
4564         @node.parent.type, if @node is not None and has a parent.
4565             OR
4566         None, otherwise.
4567     """
4568     if node is None or node.parent is None:
4569         return None
4570
4571     return node.parent.type
4572
4573
4574 def is_empty_par(leaf: Leaf) -> bool:
4575     return is_empty_lpar(leaf) or is_empty_rpar(leaf)
4576
4577
4578 def is_empty_lpar(leaf: Leaf) -> bool:
4579     return leaf.type == token.LPAR and leaf.value == ""
4580
4581
4582 def is_empty_rpar(leaf: Leaf) -> bool:
4583     return leaf.type == token.RPAR and leaf.value == ""
4584
4585
4586 def is_valid_index_factory(seq: Sequence[Any]) -> Callable[[int], bool]:
4587     """
4588     Examples:
4589         ```
4590         my_list = [1, 2, 3]
4591
4592         is_valid_index = is_valid_index_factory(my_list)
4593
4594         assert is_valid_index(0)
4595         assert is_valid_index(2)
4596
4597         assert not is_valid_index(3)
4598         assert not is_valid_index(-1)
4599         ```
4600     """
4601
4602     def is_valid_index(idx: int) -> bool:
4603         """
4604         Returns:
4605             True iff @idx is positive AND seq[@idx] does NOT raise an
4606             IndexError.
4607         """
4608         return 0 <= idx < len(seq)
4609
4610     return is_valid_index
4611
4612
4613 def line_to_string(line: Line) -> str:
4614     """Returns the string representation of @line.
4615
4616     WARNING: This is known to be computationally expensive.
4617     """
4618     return str(line).strip("\n")
4619
4620
4621 def append_leaves(
4622     new_line: Line, old_line: Line, leaves: List[Leaf], preformatted: bool = False
4623 ) -> None:
4624     """
4625     Append leaves (taken from @old_line) to @new_line, making sure to fix the
4626     underlying Node structure where appropriate.
4627
4628     All of the leaves in @leaves are duplicated. The duplicates are then
4629     appended to @new_line and used to replace their originals in the underlying
4630     Node structure. Any comments attached to the old leaves are reattached to
4631     the new leaves.
4632
4633     Pre-conditions:
4634         set(@leaves) is a subset of set(@old_line.leaves).
4635     """
4636     for old_leaf in leaves:
4637         new_leaf = Leaf(old_leaf.type, old_leaf.value)
4638         replace_child(old_leaf, new_leaf)
4639         new_line.append(new_leaf, preformatted=preformatted)
4640
4641         for comment_leaf in old_line.comments_after(old_leaf):
4642             new_line.append(comment_leaf, preformatted=True)
4643
4644
4645 def replace_child(old_child: LN, new_child: LN) -> None:
4646     """
4647     Side Effects:
4648         * If @old_child.parent is set, replace @old_child with @new_child in
4649         @old_child's underlying Node structure.
4650             OR
4651         * Otherwise, this function does nothing.
4652     """
4653     parent = old_child.parent
4654     if not parent:
4655         return
4656
4657     child_idx = old_child.remove()
4658     if child_idx is not None:
4659         parent.insert_child(child_idx, new_child)
4660
4661
4662 def get_string_prefix(string: str) -> str:
4663     """
4664     Pre-conditions:
4665         * assert_is_leaf_string(@string)
4666
4667     Returns:
4668         @string's prefix (e.g. '', 'r', 'f', or 'rf').
4669     """
4670     assert_is_leaf_string(string)
4671
4672     prefix = ""
4673     prefix_idx = 0
4674     while string[prefix_idx] in STRING_PREFIX_CHARS:
4675         prefix += string[prefix_idx].lower()
4676         prefix_idx += 1
4677
4678     return prefix
4679
4680
4681 def assert_is_leaf_string(string: str) -> None:
4682     """
4683     Checks the pre-condition that @string has the format that you would expect
4684     of `leaf.value` where `leaf` is some Leaf such that `leaf.type ==
4685     token.STRING`. A more precise description of the pre-conditions that are
4686     checked are listed below.
4687
4688     Pre-conditions:
4689         * @string starts with either ', ", <prefix>', or <prefix>" where
4690         `set(<prefix>)` is some subset of `set(STRING_PREFIX_CHARS)`.
4691         * @string ends with a quote character (' or ").
4692
4693     Raises:
4694         AssertionError(...) if the pre-conditions listed above are not
4695         satisfied.
4696     """
4697     dquote_idx = string.find('"')
4698     squote_idx = string.find("'")
4699     if -1 in [dquote_idx, squote_idx]:
4700         quote_idx = max(dquote_idx, squote_idx)
4701     else:
4702         quote_idx = min(squote_idx, dquote_idx)
4703
4704     assert (
4705         0 <= quote_idx < len(string) - 1
4706     ), f"{string!r} is missing a starting quote character (' or \")."
4707     assert string[-1] in (
4708         "'",
4709         '"',
4710     ), f"{string!r} is missing an ending quote character (' or \")."
4711     assert set(string[:quote_idx]).issubset(
4712         set(STRING_PREFIX_CHARS)
4713     ), f"{set(string[:quote_idx])} is NOT a subset of {set(STRING_PREFIX_CHARS)}."
4714
4715
4716 def left_hand_split(line: Line, _features: Collection[Feature] = ()) -> Iterator[Line]:
4717     """Split line into many lines, starting with the first matching bracket pair.
4718
4719     Note: this usually looks weird, only use this for function definitions.
4720     Prefer RHS otherwise.  This is why this function is not symmetrical with
4721     :func:`right_hand_split` which also handles optional parentheses.
4722     """
4723     tail_leaves: List[Leaf] = []
4724     body_leaves: List[Leaf] = []
4725     head_leaves: List[Leaf] = []
4726     current_leaves = head_leaves
4727     matching_bracket: Optional[Leaf] = None
4728     for leaf in line.leaves:
4729         if (
4730             current_leaves is body_leaves
4731             and leaf.type in CLOSING_BRACKETS
4732             and leaf.opening_bracket is matching_bracket
4733         ):
4734             current_leaves = tail_leaves if body_leaves else head_leaves
4735         current_leaves.append(leaf)
4736         if current_leaves is head_leaves:
4737             if leaf.type in OPENING_BRACKETS:
4738                 matching_bracket = leaf
4739                 current_leaves = body_leaves
4740     if not matching_bracket:
4741         raise CannotSplit("No brackets found")
4742
4743     head = bracket_split_build_line(head_leaves, line, matching_bracket)
4744     body = bracket_split_build_line(body_leaves, line, matching_bracket, is_body=True)
4745     tail = bracket_split_build_line(tail_leaves, line, matching_bracket)
4746     bracket_split_succeeded_or_raise(head, body, tail)
4747     for result in (head, body, tail):
4748         if result:
4749             yield result
4750
4751
4752 def right_hand_split(
4753     line: Line,
4754     line_length: int,
4755     features: Collection[Feature] = (),
4756     omit: Collection[LeafID] = (),
4757 ) -> Iterator[Line]:
4758     """Split line into many lines, starting with the last matching bracket pair.
4759
4760     If the split was by optional parentheses, attempt splitting without them, too.
4761     `omit` is a collection of closing bracket IDs that shouldn't be considered for
4762     this split.
4763
4764     Note: running this function modifies `bracket_depth` on the leaves of `line`.
4765     """
4766     tail_leaves: List[Leaf] = []
4767     body_leaves: List[Leaf] = []
4768     head_leaves: List[Leaf] = []
4769     current_leaves = tail_leaves
4770     opening_bracket: Optional[Leaf] = None
4771     closing_bracket: Optional[Leaf] = None
4772     for leaf in reversed(line.leaves):
4773         if current_leaves is body_leaves:
4774             if leaf is opening_bracket:
4775                 current_leaves = head_leaves if body_leaves else tail_leaves
4776         current_leaves.append(leaf)
4777         if current_leaves is tail_leaves:
4778             if leaf.type in CLOSING_BRACKETS and id(leaf) not in omit:
4779                 opening_bracket = leaf.opening_bracket
4780                 closing_bracket = leaf
4781                 current_leaves = body_leaves
4782     if not (opening_bracket and closing_bracket and head_leaves):
4783         # If there is no opening or closing_bracket that means the split failed and
4784         # all content is in the tail.  Otherwise, if `head_leaves` are empty, it means
4785         # the matching `opening_bracket` wasn't available on `line` anymore.
4786         raise CannotSplit("No brackets found")
4787
4788     tail_leaves.reverse()
4789     body_leaves.reverse()
4790     head_leaves.reverse()
4791     head = bracket_split_build_line(head_leaves, line, opening_bracket)
4792     body = bracket_split_build_line(body_leaves, line, opening_bracket, is_body=True)
4793     tail = bracket_split_build_line(tail_leaves, line, opening_bracket)
4794     bracket_split_succeeded_or_raise(head, body, tail)
4795     if (
4796         Feature.FORCE_OPTIONAL_PARENTHESES not in features
4797         # the opening bracket is an optional paren
4798         and opening_bracket.type == token.LPAR
4799         and not opening_bracket.value
4800         # the closing bracket is an optional paren
4801         and closing_bracket.type == token.RPAR
4802         and not closing_bracket.value
4803         # it's not an import (optional parens are the only thing we can split on
4804         # in this case; attempting a split without them is a waste of time)
4805         and not line.is_import
4806         # there are no standalone comments in the body
4807         and not body.contains_standalone_comments(0)
4808         # and we can actually remove the parens
4809         and can_omit_invisible_parens(body, line_length, omit_on_explode=omit)
4810     ):
4811         omit = {id(closing_bracket), *omit}
4812         try:
4813             yield from right_hand_split(line, line_length, features=features, omit=omit)
4814             return
4815
4816         except CannotSplit:
4817             if not (
4818                 can_be_split(body)
4819                 or is_line_short_enough(body, line_length=line_length)
4820             ):
4821                 raise CannotSplit(
4822                     "Splitting failed, body is still too long and can't be split."
4823                 )
4824
4825             elif head.contains_multiline_strings() or tail.contains_multiline_strings():
4826                 raise CannotSplit(
4827                     "The current optional pair of parentheses is bound to fail to"
4828                     " satisfy the splitting algorithm because the head or the tail"
4829                     " contains multiline strings which by definition never fit one"
4830                     " line."
4831                 )
4832
4833     ensure_visible(opening_bracket)
4834     ensure_visible(closing_bracket)
4835     for result in (head, body, tail):
4836         if result:
4837             yield result
4838
4839
4840 def bracket_split_succeeded_or_raise(head: Line, body: Line, tail: Line) -> None:
4841     """Raise :exc:`CannotSplit` if the last left- or right-hand split failed.
4842
4843     Do nothing otherwise.
4844
4845     A left- or right-hand split is based on a pair of brackets. Content before
4846     (and including) the opening bracket is left on one line, content inside the
4847     brackets is put on a separate line, and finally content starting with and
4848     following the closing bracket is put on a separate line.
4849
4850     Those are called `head`, `body`, and `tail`, respectively. If the split
4851     produced the same line (all content in `head`) or ended up with an empty `body`
4852     and the `tail` is just the closing bracket, then it's considered failed.
4853     """
4854     tail_len = len(str(tail).strip())
4855     if not body:
4856         if tail_len == 0:
4857             raise CannotSplit("Splitting brackets produced the same line")
4858
4859         elif tail_len < 3:
4860             raise CannotSplit(
4861                 f"Splitting brackets on an empty body to save {tail_len} characters is"
4862                 " not worth it"
4863             )
4864
4865
4866 def bracket_split_build_line(
4867     leaves: List[Leaf], original: Line, opening_bracket: Leaf, *, is_body: bool = False
4868 ) -> Line:
4869     """Return a new line with given `leaves` and respective comments from `original`.
4870
4871     If `is_body` is True, the result line is one-indented inside brackets and as such
4872     has its first leaf's prefix normalized and a trailing comma added when expected.
4873     """
4874     result = Line(depth=original.depth)
4875     if is_body:
4876         result.inside_brackets = True
4877         result.depth += 1
4878         if leaves:
4879             # Since body is a new indent level, remove spurious leading whitespace.
4880             normalize_prefix(leaves[0], inside_brackets=True)
4881             # Ensure a trailing comma for imports and standalone function arguments, but
4882             # be careful not to add one after any comments or within type annotations.
4883             no_commas = (
4884                 original.is_def
4885                 and opening_bracket.value == "("
4886                 and not any(leaf.type == token.COMMA for leaf in leaves)
4887             )
4888
4889             if original.is_import or no_commas:
4890                 for i in range(len(leaves) - 1, -1, -1):
4891                     if leaves[i].type == STANDALONE_COMMENT:
4892                         continue
4893
4894                     if leaves[i].type != token.COMMA:
4895                         new_comma = Leaf(token.COMMA, ",")
4896                         leaves.insert(i + 1, new_comma)
4897                     break
4898
4899     # Populate the line
4900     for leaf in leaves:
4901         result.append(leaf, preformatted=True)
4902         for comment_after in original.comments_after(leaf):
4903             result.append(comment_after, preformatted=True)
4904     if is_body and should_split_body_explode(result, opening_bracket):
4905         result.should_explode = True
4906     return result
4907
4908
4909 def dont_increase_indentation(split_func: Transformer) -> Transformer:
4910     """Normalize prefix of the first leaf in every line returned by `split_func`.
4911
4912     This is a decorator over relevant split functions.
4913     """
4914
4915     @wraps(split_func)
4916     def split_wrapper(line: Line, features: Collection[Feature] = ()) -> Iterator[Line]:
4917         for line in split_func(line, features):
4918             normalize_prefix(line.leaves[0], inside_brackets=True)
4919             yield line
4920
4921     return split_wrapper
4922
4923
4924 @dont_increase_indentation
4925 def delimiter_split(line: Line, features: Collection[Feature] = ()) -> Iterator[Line]:
4926     """Split according to delimiters of the highest priority.
4927
4928     If the appropriate Features are given, the split will add trailing commas
4929     also in function signatures and calls that contain `*` and `**`.
4930     """
4931     try:
4932         last_leaf = line.leaves[-1]
4933     except IndexError:
4934         raise CannotSplit("Line empty")
4935
4936     bt = line.bracket_tracker
4937     try:
4938         delimiter_priority = bt.max_delimiter_priority(exclude={id(last_leaf)})
4939     except ValueError:
4940         raise CannotSplit("No delimiters found")
4941
4942     if delimiter_priority == DOT_PRIORITY:
4943         if bt.delimiter_count_with_priority(delimiter_priority) == 1:
4944             raise CannotSplit("Splitting a single attribute from its owner looks wrong")
4945
4946     current_line = Line(depth=line.depth, inside_brackets=line.inside_brackets)
4947     lowest_depth = sys.maxsize
4948     trailing_comma_safe = True
4949
4950     def append_to_line(leaf: Leaf) -> Iterator[Line]:
4951         """Append `leaf` to current line or to new line if appending impossible."""
4952         nonlocal current_line
4953         try:
4954             current_line.append_safe(leaf, preformatted=True)
4955         except ValueError:
4956             yield current_line
4957
4958             current_line = Line(depth=line.depth, inside_brackets=line.inside_brackets)
4959             current_line.append(leaf)
4960
4961     for leaf in line.leaves:
4962         yield from append_to_line(leaf)
4963
4964         for comment_after in line.comments_after(leaf):
4965             yield from append_to_line(comment_after)
4966
4967         lowest_depth = min(lowest_depth, leaf.bracket_depth)
4968         if leaf.bracket_depth == lowest_depth:
4969             if is_vararg(leaf, within={syms.typedargslist}):
4970                 trailing_comma_safe = (
4971                     trailing_comma_safe and Feature.TRAILING_COMMA_IN_DEF in features
4972                 )
4973             elif is_vararg(leaf, within={syms.arglist, syms.argument}):
4974                 trailing_comma_safe = (
4975                     trailing_comma_safe and Feature.TRAILING_COMMA_IN_CALL in features
4976                 )
4977
4978         leaf_priority = bt.delimiters.get(id(leaf))
4979         if leaf_priority == delimiter_priority:
4980             yield current_line
4981
4982             current_line = Line(depth=line.depth, inside_brackets=line.inside_brackets)
4983     if current_line:
4984         if (
4985             trailing_comma_safe
4986             and delimiter_priority == COMMA_PRIORITY
4987             and current_line.leaves[-1].type != token.COMMA
4988             and current_line.leaves[-1].type != STANDALONE_COMMENT
4989         ):
4990             new_comma = Leaf(token.COMMA, ",")
4991             current_line.append(new_comma)
4992         yield current_line
4993
4994
4995 @dont_increase_indentation
4996 def standalone_comment_split(
4997     line: Line, features: Collection[Feature] = ()
4998 ) -> Iterator[Line]:
4999     """Split standalone comments from the rest of the line."""
5000     if not line.contains_standalone_comments(0):
5001         raise CannotSplit("Line does not have any standalone comments")
5002
5003     current_line = Line(depth=line.depth, inside_brackets=line.inside_brackets)
5004
5005     def append_to_line(leaf: Leaf) -> Iterator[Line]:
5006         """Append `leaf` to current line or to new line if appending impossible."""
5007         nonlocal current_line
5008         try:
5009             current_line.append_safe(leaf, preformatted=True)
5010         except ValueError:
5011             yield current_line
5012
5013             current_line = Line(depth=line.depth, inside_brackets=line.inside_brackets)
5014             current_line.append(leaf)
5015
5016     for leaf in line.leaves:
5017         yield from append_to_line(leaf)
5018
5019         for comment_after in line.comments_after(leaf):
5020             yield from append_to_line(comment_after)
5021
5022     if current_line:
5023         yield current_line
5024
5025
5026 def is_import(leaf: Leaf) -> bool:
5027     """Return True if the given leaf starts an import statement."""
5028     p = leaf.parent
5029     t = leaf.type
5030     v = leaf.value
5031     return bool(
5032         t == token.NAME
5033         and (
5034             (v == "import" and p and p.type == syms.import_name)
5035             or (v == "from" and p and p.type == syms.import_from)
5036         )
5037     )
5038
5039
5040 def is_type_comment(leaf: Leaf, suffix: str = "") -> bool:
5041     """Return True if the given leaf is a special comment.
5042     Only returns true for type comments for now."""
5043     t = leaf.type
5044     v = leaf.value
5045     return t in {token.COMMENT, STANDALONE_COMMENT} and v.startswith("# type:" + suffix)
5046
5047
5048 def normalize_prefix(leaf: Leaf, *, inside_brackets: bool) -> None:
5049     """Leave existing extra newlines if not `inside_brackets`. Remove everything
5050     else.
5051
5052     Note: don't use backslashes for formatting or you'll lose your voting rights.
5053     """
5054     if not inside_brackets:
5055         spl = leaf.prefix.split("#")
5056         if "\\" not in spl[0]:
5057             nl_count = spl[-1].count("\n")
5058             if len(spl) > 1:
5059                 nl_count -= 1
5060             leaf.prefix = "\n" * nl_count
5061             return
5062
5063     leaf.prefix = ""
5064
5065
5066 def normalize_string_prefix(leaf: Leaf, remove_u_prefix: bool = False) -> None:
5067     """Make all string prefixes lowercase.
5068
5069     If remove_u_prefix is given, also removes any u prefix from the string.
5070
5071     Note: Mutates its argument.
5072     """
5073     match = re.match(r"^([" + STRING_PREFIX_CHARS + r"]*)(.*)$", leaf.value, re.DOTALL)
5074     assert match is not None, f"failed to match string {leaf.value!r}"
5075     orig_prefix = match.group(1)
5076     new_prefix = orig_prefix.replace("F", "f").replace("B", "b").replace("U", "u")
5077     if remove_u_prefix:
5078         new_prefix = new_prefix.replace("u", "")
5079     leaf.value = f"{new_prefix}{match.group(2)}"
5080
5081
5082 def normalize_string_quotes(leaf: Leaf) -> None:
5083     """Prefer double quotes but only if it doesn't cause more escaping.
5084
5085     Adds or removes backslashes as appropriate. Doesn't parse and fix
5086     strings nested in f-strings (yet).
5087
5088     Note: Mutates its argument.
5089     """
5090     value = leaf.value.lstrip(STRING_PREFIX_CHARS)
5091     if value[:3] == '"""':
5092         return
5093
5094     elif value[:3] == "'''":
5095         orig_quote = "'''"
5096         new_quote = '"""'
5097     elif value[0] == '"':
5098         orig_quote = '"'
5099         new_quote = "'"
5100     else:
5101         orig_quote = "'"
5102         new_quote = '"'
5103     first_quote_pos = leaf.value.find(orig_quote)
5104     if first_quote_pos == -1:
5105         return  # There's an internal error
5106
5107     prefix = leaf.value[:first_quote_pos]
5108     unescaped_new_quote = re.compile(rf"(([^\\]|^)(\\\\)*){new_quote}")
5109     escaped_new_quote = re.compile(rf"([^\\]|^)\\((?:\\\\)*){new_quote}")
5110     escaped_orig_quote = re.compile(rf"([^\\]|^)\\((?:\\\\)*){orig_quote}")
5111     body = leaf.value[first_quote_pos + len(orig_quote) : -len(orig_quote)]
5112     if "r" in prefix.casefold():
5113         if unescaped_new_quote.search(body):
5114             # There's at least one unescaped new_quote in this raw string
5115             # so converting is impossible
5116             return
5117
5118         # Do not introduce or remove backslashes in raw strings
5119         new_body = body
5120     else:
5121         # remove unnecessary escapes
5122         new_body = sub_twice(escaped_new_quote, rf"\1\2{new_quote}", body)
5123         if body != new_body:
5124             # Consider the string without unnecessary escapes as the original
5125             body = new_body
5126             leaf.value = f"{prefix}{orig_quote}{body}{orig_quote}"
5127         new_body = sub_twice(escaped_orig_quote, rf"\1\2{orig_quote}", new_body)
5128         new_body = sub_twice(unescaped_new_quote, rf"\1\\{new_quote}", new_body)
5129     if "f" in prefix.casefold():
5130         matches = re.findall(
5131             r"""
5132             (?:[^{]|^)\{  # start of the string or a non-{ followed by a single {
5133                 ([^{].*?)  # contents of the brackets except if begins with {{
5134             \}(?:[^}]|$)  # A } followed by end of the string or a non-}
5135             """,
5136             new_body,
5137             re.VERBOSE,
5138         )
5139         for m in matches:
5140             if "\\" in str(m):
5141                 # Do not introduce backslashes in interpolated expressions
5142                 return
5143
5144     if new_quote == '"""' and new_body[-1:] == '"':
5145         # edge case:
5146         new_body = new_body[:-1] + '\\"'
5147     orig_escape_count = body.count("\\")
5148     new_escape_count = new_body.count("\\")
5149     if new_escape_count > orig_escape_count:
5150         return  # Do not introduce more escaping
5151
5152     if new_escape_count == orig_escape_count and orig_quote == '"':
5153         return  # Prefer double quotes
5154
5155     leaf.value = f"{prefix}{new_quote}{new_body}{new_quote}"
5156
5157
5158 def normalize_numeric_literal(leaf: Leaf) -> None:
5159     """Normalizes numeric (float, int, and complex) literals.
5160
5161     All letters used in the representation are normalized to lowercase (except
5162     in Python 2 long literals).
5163     """
5164     text = leaf.value.lower()
5165     if text.startswith(("0o", "0b")):
5166         # Leave octal and binary literals alone.
5167         pass
5168     elif text.startswith("0x"):
5169         # Change hex literals to upper case.
5170         before, after = text[:2], text[2:]
5171         text = f"{before}{after.upper()}"
5172     elif "e" in text:
5173         before, after = text.split("e")
5174         sign = ""
5175         if after.startswith("-"):
5176             after = after[1:]
5177             sign = "-"
5178         elif after.startswith("+"):
5179             after = after[1:]
5180         before = format_float_or_int_string(before)
5181         text = f"{before}e{sign}{after}"
5182     elif text.endswith(("j", "l")):
5183         number = text[:-1]
5184         suffix = text[-1]
5185         # Capitalize in "2L" because "l" looks too similar to "1".
5186         if suffix == "l":
5187             suffix = "L"
5188         text = f"{format_float_or_int_string(number)}{suffix}"
5189     else:
5190         text = format_float_or_int_string(text)
5191     leaf.value = text
5192
5193
5194 def format_float_or_int_string(text: str) -> str:
5195     """Formats a float string like "1.0"."""
5196     if "." not in text:
5197         return text
5198
5199     before, after = text.split(".")
5200     return f"{before or 0}.{after or 0}"
5201
5202
5203 def normalize_invisible_parens(node: Node, parens_after: Set[str]) -> None:
5204     """Make existing optional parentheses invisible or create new ones.
5205
5206     `parens_after` is a set of string leaf values immediately after which parens
5207     should be put.
5208
5209     Standardizes on visible parentheses for single-element tuples, and keeps
5210     existing visible parentheses for other tuples and generator expressions.
5211     """
5212     for pc in list_comments(node.prefix, is_endmarker=False):
5213         if pc.value in FMT_OFF:
5214             # This `node` has a prefix with `# fmt: off`, don't mess with parens.
5215             return
5216     check_lpar = False
5217     for index, child in enumerate(list(node.children)):
5218         # Fixes a bug where invisible parens are not properly stripped from
5219         # assignment statements that contain type annotations.
5220         if isinstance(child, Node) and child.type == syms.annassign:
5221             normalize_invisible_parens(child, parens_after=parens_after)
5222
5223         # Add parentheses around long tuple unpacking in assignments.
5224         if (
5225             index == 0
5226             and isinstance(child, Node)
5227             and child.type == syms.testlist_star_expr
5228         ):
5229             check_lpar = True
5230
5231         if check_lpar:
5232             if is_walrus_assignment(child):
5233                 pass
5234
5235             elif child.type == syms.atom:
5236                 if maybe_make_parens_invisible_in_atom(child, parent=node):
5237                     wrap_in_parentheses(node, child, visible=False)
5238             elif is_one_tuple(child):
5239                 wrap_in_parentheses(node, child, visible=True)
5240             elif node.type == syms.import_from:
5241                 # "import from" nodes store parentheses directly as part of
5242                 # the statement
5243                 if child.type == token.LPAR:
5244                     # make parentheses invisible
5245                     child.value = ""  # type: ignore
5246                     node.children[-1].value = ""  # type: ignore
5247                 elif child.type != token.STAR:
5248                     # insert invisible parentheses
5249                     node.insert_child(index, Leaf(token.LPAR, ""))
5250                     node.append_child(Leaf(token.RPAR, ""))
5251                 break
5252
5253             elif not (isinstance(child, Leaf) and is_multiline_string(child)):
5254                 wrap_in_parentheses(node, child, visible=False)
5255
5256         check_lpar = isinstance(child, Leaf) and child.value in parens_after
5257
5258
5259 def normalize_fmt_off(node: Node) -> None:
5260     """Convert content between `# fmt: off`/`# fmt: on` into standalone comments."""
5261     try_again = True
5262     while try_again:
5263         try_again = convert_one_fmt_off_pair(node)
5264
5265
5266 def convert_one_fmt_off_pair(node: Node) -> bool:
5267     """Convert content of a single `# fmt: off`/`# fmt: on` into a standalone comment.
5268
5269     Returns True if a pair was converted.
5270     """
5271     for leaf in node.leaves():
5272         previous_consumed = 0
5273         for comment in list_comments(leaf.prefix, is_endmarker=False):
5274             if comment.value in FMT_OFF:
5275                 # We only want standalone comments. If there's no previous leaf or
5276                 # the previous leaf is indentation, it's a standalone comment in
5277                 # disguise.
5278                 if comment.type != STANDALONE_COMMENT:
5279                     prev = preceding_leaf(leaf)
5280                     if prev and prev.type not in WHITESPACE:
5281                         continue
5282
5283                 ignored_nodes = list(generate_ignored_nodes(leaf))
5284                 if not ignored_nodes:
5285                     continue
5286
5287                 first = ignored_nodes[0]  # Can be a container node with the `leaf`.
5288                 parent = first.parent
5289                 prefix = first.prefix
5290                 first.prefix = prefix[comment.consumed :]
5291                 hidden_value = (
5292                     comment.value + "\n" + "".join(str(n) for n in ignored_nodes)
5293                 )
5294                 if hidden_value.endswith("\n"):
5295                     # That happens when one of the `ignored_nodes` ended with a NEWLINE
5296                     # leaf (possibly followed by a DEDENT).
5297                     hidden_value = hidden_value[:-1]
5298                 first_idx: Optional[int] = None
5299                 for ignored in ignored_nodes:
5300                     index = ignored.remove()
5301                     if first_idx is None:
5302                         first_idx = index
5303                 assert parent is not None, "INTERNAL ERROR: fmt: on/off handling (1)"
5304                 assert first_idx is not None, "INTERNAL ERROR: fmt: on/off handling (2)"
5305                 parent.insert_child(
5306                     first_idx,
5307                     Leaf(
5308                         STANDALONE_COMMENT,
5309                         hidden_value,
5310                         prefix=prefix[:previous_consumed] + "\n" * comment.newlines,
5311                     ),
5312                 )
5313                 return True
5314
5315             previous_consumed = comment.consumed
5316
5317     return False
5318
5319
5320 def generate_ignored_nodes(leaf: Leaf) -> Iterator[LN]:
5321     """Starting from the container of `leaf`, generate all leaves until `# fmt: on`.
5322
5323     Stops at the end of the block.
5324     """
5325     container: Optional[LN] = container_of(leaf)
5326     while container is not None and container.type != token.ENDMARKER:
5327         if is_fmt_on(container):
5328             return
5329
5330         # fix for fmt: on in children
5331         if contains_fmt_on_at_column(container, leaf.column):
5332             for child in container.children:
5333                 if contains_fmt_on_at_column(child, leaf.column):
5334                     return
5335                 yield child
5336         else:
5337             yield container
5338             container = container.next_sibling
5339
5340
5341 def is_fmt_on(container: LN) -> bool:
5342     """Determine whether formatting is switched on within a container.
5343     Determined by whether the last `# fmt:` comment is `on` or `off`.
5344     """
5345     fmt_on = False
5346     for comment in list_comments(container.prefix, is_endmarker=False):
5347         if comment.value in FMT_ON:
5348             fmt_on = True
5349         elif comment.value in FMT_OFF:
5350             fmt_on = False
5351     return fmt_on
5352
5353
5354 def contains_fmt_on_at_column(container: LN, column: int) -> bool:
5355     """Determine if children at a given column have formatting switched on."""
5356     for child in container.children:
5357         if (
5358             isinstance(child, Node)
5359             and first_leaf_column(child) == column
5360             or isinstance(child, Leaf)
5361             and child.column == column
5362         ):
5363             if is_fmt_on(child):
5364                 return True
5365
5366     return False
5367
5368
5369 def first_leaf_column(node: Node) -> Optional[int]:
5370     """Returns the column of the first leaf child of a node."""
5371     for child in node.children:
5372         if isinstance(child, Leaf):
5373             return child.column
5374     return None
5375
5376
5377 def maybe_make_parens_invisible_in_atom(node: LN, parent: LN) -> bool:
5378     """If it's safe, make the parens in the atom `node` invisible, recursively.
5379     Additionally, remove repeated, adjacent invisible parens from the atom `node`
5380     as they are redundant.
5381
5382     Returns whether the node should itself be wrapped in invisible parentheses.
5383
5384     """
5385     if (
5386         node.type != syms.atom
5387         or is_empty_tuple(node)
5388         or is_one_tuple(node)
5389         or (is_yield(node) and parent.type != syms.expr_stmt)
5390         or max_delimiter_priority_in_atom(node) >= COMMA_PRIORITY
5391     ):
5392         return False
5393
5394     first = node.children[0]
5395     last = node.children[-1]
5396     if first.type == token.LPAR and last.type == token.RPAR:
5397         middle = node.children[1]
5398         # make parentheses invisible
5399         first.value = ""  # type: ignore
5400         last.value = ""  # type: ignore
5401         maybe_make_parens_invisible_in_atom(middle, parent=parent)
5402
5403         if is_atom_with_invisible_parens(middle):
5404             # Strip the invisible parens from `middle` by replacing
5405             # it with the child in-between the invisible parens
5406             middle.replace(middle.children[1])
5407
5408         return False
5409
5410     return True
5411
5412
5413 def is_atom_with_invisible_parens(node: LN) -> bool:
5414     """Given a `LN`, determines whether it's an atom `node` with invisible
5415     parens. Useful in dedupe-ing and normalizing parens.
5416     """
5417     if isinstance(node, Leaf) or node.type != syms.atom:
5418         return False
5419
5420     first, last = node.children[0], node.children[-1]
5421     return (
5422         isinstance(first, Leaf)
5423         and first.type == token.LPAR
5424         and first.value == ""
5425         and isinstance(last, Leaf)
5426         and last.type == token.RPAR
5427         and last.value == ""
5428     )
5429
5430
5431 def is_empty_tuple(node: LN) -> bool:
5432     """Return True if `node` holds an empty tuple."""
5433     return (
5434         node.type == syms.atom
5435         and len(node.children) == 2
5436         and node.children[0].type == token.LPAR
5437         and node.children[1].type == token.RPAR
5438     )
5439
5440
5441 def unwrap_singleton_parenthesis(node: LN) -> Optional[LN]:
5442     """Returns `wrapped` if `node` is of the shape ( wrapped ).
5443
5444     Parenthesis can be optional. Returns None otherwise"""
5445     if len(node.children) != 3:
5446         return None
5447
5448     lpar, wrapped, rpar = node.children
5449     if not (lpar.type == token.LPAR and rpar.type == token.RPAR):
5450         return None
5451
5452     return wrapped
5453
5454
5455 def wrap_in_parentheses(parent: Node, child: LN, *, visible: bool = True) -> None:
5456     """Wrap `child` in parentheses.
5457
5458     This replaces `child` with an atom holding the parentheses and the old
5459     child.  That requires moving the prefix.
5460
5461     If `visible` is False, the leaves will be valueless (and thus invisible).
5462     """
5463     lpar = Leaf(token.LPAR, "(" if visible else "")
5464     rpar = Leaf(token.RPAR, ")" if visible else "")
5465     prefix = child.prefix
5466     child.prefix = ""
5467     index = child.remove() or 0
5468     new_child = Node(syms.atom, [lpar, child, rpar])
5469     new_child.prefix = prefix
5470     parent.insert_child(index, new_child)
5471
5472
5473 def is_one_tuple(node: LN) -> bool:
5474     """Return True if `node` holds a tuple with one element, with or without parens."""
5475     if node.type == syms.atom:
5476         gexp = unwrap_singleton_parenthesis(node)
5477         if gexp is None or gexp.type != syms.testlist_gexp:
5478             return False
5479
5480         return len(gexp.children) == 2 and gexp.children[1].type == token.COMMA
5481
5482     return (
5483         node.type in IMPLICIT_TUPLE
5484         and len(node.children) == 2
5485         and node.children[1].type == token.COMMA
5486     )
5487
5488
5489 def is_walrus_assignment(node: LN) -> bool:
5490     """Return True iff `node` is of the shape ( test := test )"""
5491     inner = unwrap_singleton_parenthesis(node)
5492     return inner is not None and inner.type == syms.namedexpr_test
5493
5494
5495 def is_yield(node: LN) -> bool:
5496     """Return True if `node` holds a `yield` or `yield from` expression."""
5497     if node.type == syms.yield_expr:
5498         return True
5499
5500     if node.type == token.NAME and node.value == "yield":  # type: ignore
5501         return True
5502
5503     if node.type != syms.atom:
5504         return False
5505
5506     if len(node.children) != 3:
5507         return False
5508
5509     lpar, expr, rpar = node.children
5510     if lpar.type == token.LPAR and rpar.type == token.RPAR:
5511         return is_yield(expr)
5512
5513     return False
5514
5515
5516 def is_vararg(leaf: Leaf, within: Set[NodeType]) -> bool:
5517     """Return True if `leaf` is a star or double star in a vararg or kwarg.
5518
5519     If `within` includes VARARGS_PARENTS, this applies to function signatures.
5520     If `within` includes UNPACKING_PARENTS, it applies to right hand-side
5521     extended iterable unpacking (PEP 3132) and additional unpacking
5522     generalizations (PEP 448).
5523     """
5524     if leaf.type not in VARARGS_SPECIALS or not leaf.parent:
5525         return False
5526
5527     p = leaf.parent
5528     if p.type == syms.star_expr:
5529         # Star expressions are also used as assignment targets in extended
5530         # iterable unpacking (PEP 3132).  See what its parent is instead.
5531         if not p.parent:
5532             return False
5533
5534         p = p.parent
5535
5536     return p.type in within
5537
5538
5539 def is_multiline_string(leaf: Leaf) -> bool:
5540     """Return True if `leaf` is a multiline string that actually spans many lines."""
5541     return has_triple_quotes(leaf.value) and "\n" in leaf.value
5542
5543
5544 def is_stub_suite(node: Node) -> bool:
5545     """Return True if `node` is a suite with a stub body."""
5546     if (
5547         len(node.children) != 4
5548         or node.children[0].type != token.NEWLINE
5549         or node.children[1].type != token.INDENT
5550         or node.children[3].type != token.DEDENT
5551     ):
5552         return False
5553
5554     return is_stub_body(node.children[2])
5555
5556
5557 def is_stub_body(node: LN) -> bool:
5558     """Return True if `node` is a simple statement containing an ellipsis."""
5559     if not isinstance(node, Node) or node.type != syms.simple_stmt:
5560         return False
5561
5562     if len(node.children) != 2:
5563         return False
5564
5565     child = node.children[0]
5566     return (
5567         child.type == syms.atom
5568         and len(child.children) == 3
5569         and all(leaf == Leaf(token.DOT, ".") for leaf in child.children)
5570     )
5571
5572
5573 def max_delimiter_priority_in_atom(node: LN) -> Priority:
5574     """Return maximum delimiter priority inside `node`.
5575
5576     This is specific to atoms with contents contained in a pair of parentheses.
5577     If `node` isn't an atom or there are no enclosing parentheses, returns 0.
5578     """
5579     if node.type != syms.atom:
5580         return 0
5581
5582     first = node.children[0]
5583     last = node.children[-1]
5584     if not (first.type == token.LPAR and last.type == token.RPAR):
5585         return 0
5586
5587     bt = BracketTracker()
5588     for c in node.children[1:-1]:
5589         if isinstance(c, Leaf):
5590             bt.mark(c)
5591         else:
5592             for leaf in c.leaves():
5593                 bt.mark(leaf)
5594     try:
5595         return bt.max_delimiter_priority()
5596
5597     except ValueError:
5598         return 0
5599
5600
5601 def ensure_visible(leaf: Leaf) -> None:
5602     """Make sure parentheses are visible.
5603
5604     They could be invisible as part of some statements (see
5605     :func:`normalize_invisible_parens` and :func:`visit_import_from`).
5606     """
5607     if leaf.type == token.LPAR:
5608         leaf.value = "("
5609     elif leaf.type == token.RPAR:
5610         leaf.value = ")"
5611
5612
5613 def should_split_body_explode(line: Line, opening_bracket: Leaf) -> bool:
5614     """Should `line` be immediately split with `delimiter_split()` after RHS?"""
5615
5616     if not (opening_bracket.parent and opening_bracket.value in "[{("):
5617         return False
5618
5619     # We're essentially checking if the body is delimited by commas and there's more
5620     # than one of them (we're excluding the trailing comma and if the delimiter priority
5621     # is still commas, that means there's more).
5622     exclude = set()
5623     trailing_comma = False
5624     try:
5625         last_leaf = line.leaves[-1]
5626         if last_leaf.type == token.COMMA:
5627             trailing_comma = True
5628             exclude.add(id(last_leaf))
5629         max_priority = line.bracket_tracker.max_delimiter_priority(exclude=exclude)
5630     except (IndexError, ValueError):
5631         return False
5632
5633     return max_priority == COMMA_PRIORITY and (
5634         trailing_comma
5635         # always explode imports
5636         or opening_bracket.parent.type in {syms.atom, syms.import_from}
5637     )
5638
5639
5640 def is_one_tuple_between(opening: Leaf, closing: Leaf, leaves: List[Leaf]) -> bool:
5641     """Return True if content between `opening` and `closing` looks like a one-tuple."""
5642     if opening.type != token.LPAR and closing.type != token.RPAR:
5643         return False
5644
5645     depth = closing.bracket_depth + 1
5646     for _opening_index, leaf in enumerate(leaves):
5647         if leaf is opening:
5648             break
5649
5650     else:
5651         raise LookupError("Opening paren not found in `leaves`")
5652
5653     commas = 0
5654     _opening_index += 1
5655     for leaf in leaves[_opening_index:]:
5656         if leaf is closing:
5657             break
5658
5659         bracket_depth = leaf.bracket_depth
5660         if bracket_depth == depth and leaf.type == token.COMMA:
5661             commas += 1
5662             if leaf.parent and leaf.parent.type in {
5663                 syms.arglist,
5664                 syms.typedargslist,
5665             }:
5666                 commas += 1
5667                 break
5668
5669     return commas < 2
5670
5671
5672 def get_features_used(node: Node) -> Set[Feature]:
5673     """Return a set of (relatively) new Python features used in this file.
5674
5675     Currently looking for:
5676     - f-strings;
5677     - underscores in numeric literals;
5678     - trailing commas after * or ** in function signatures and calls;
5679     - positional only arguments in function signatures and lambdas;
5680     """
5681     features: Set[Feature] = set()
5682     for n in node.pre_order():
5683         if n.type == token.STRING:
5684             value_head = n.value[:2]  # type: ignore
5685             if value_head in {'f"', 'F"', "f'", "F'", "rf", "fr", "RF", "FR"}:
5686                 features.add(Feature.F_STRINGS)
5687
5688         elif n.type == token.NUMBER:
5689             if "_" in n.value:  # type: ignore
5690                 features.add(Feature.NUMERIC_UNDERSCORES)
5691
5692         elif n.type == token.SLASH:
5693             if n.parent and n.parent.type in {syms.typedargslist, syms.arglist}:
5694                 features.add(Feature.POS_ONLY_ARGUMENTS)
5695
5696         elif n.type == token.COLONEQUAL:
5697             features.add(Feature.ASSIGNMENT_EXPRESSIONS)
5698
5699         elif (
5700             n.type in {syms.typedargslist, syms.arglist}
5701             and n.children
5702             and n.children[-1].type == token.COMMA
5703         ):
5704             if n.type == syms.typedargslist:
5705                 feature = Feature.TRAILING_COMMA_IN_DEF
5706             else:
5707                 feature = Feature.TRAILING_COMMA_IN_CALL
5708
5709             for ch in n.children:
5710                 if ch.type in STARS:
5711                     features.add(feature)
5712
5713                 if ch.type == syms.argument:
5714                     for argch in ch.children:
5715                         if argch.type in STARS:
5716                             features.add(feature)
5717
5718     return features
5719
5720
5721 def detect_target_versions(node: Node) -> Set[TargetVersion]:
5722     """Detect the version to target based on the nodes used."""
5723     features = get_features_used(node)
5724     return {
5725         version for version in TargetVersion if features <= VERSION_TO_FEATURES[version]
5726     }
5727
5728
5729 def generate_trailers_to_omit(line: Line, line_length: int) -> Iterator[Set[LeafID]]:
5730     """Generate sets of closing bracket IDs that should be omitted in a RHS.
5731
5732     Brackets can be omitted if the entire trailer up to and including
5733     a preceding closing bracket fits in one line.
5734
5735     Yielded sets are cumulative (contain results of previous yields, too).  First
5736     set is empty, unless the line should explode, in which case bracket pairs until
5737     the one that needs to explode are omitted.
5738     """
5739
5740     omit: Set[LeafID] = set()
5741     if not line.should_explode:
5742         yield omit
5743
5744     length = 4 * line.depth
5745     opening_bracket: Optional[Leaf] = None
5746     closing_bracket: Optional[Leaf] = None
5747     inner_brackets: Set[LeafID] = set()
5748     for index, leaf, leaf_length in enumerate_with_length(line, reversed=True):
5749         length += leaf_length
5750         if length > line_length:
5751             break
5752
5753         has_inline_comment = leaf_length > len(leaf.value) + len(leaf.prefix)
5754         if leaf.type == STANDALONE_COMMENT or has_inline_comment:
5755             break
5756
5757         if opening_bracket:
5758             if leaf is opening_bracket:
5759                 opening_bracket = None
5760             elif leaf.type in CLOSING_BRACKETS:
5761                 prev = line.leaves[index - 1] if index > 0 else None
5762                 if (
5763                     line.should_explode
5764                     and prev
5765                     and prev.type == token.COMMA
5766                     and not is_one_tuple_between(
5767                         leaf.opening_bracket, leaf, line.leaves
5768                     )
5769                 ):
5770                     # Never omit bracket pairs with trailing commas.
5771                     # We need to explode on those.
5772                     break
5773
5774                 inner_brackets.add(id(leaf))
5775         elif leaf.type in CLOSING_BRACKETS:
5776             prev = line.leaves[index - 1] if index > 0 else None
5777             if prev and prev.type in OPENING_BRACKETS:
5778                 # Empty brackets would fail a split so treat them as "inner"
5779                 # brackets (e.g. only add them to the `omit` set if another
5780                 # pair of brackets was good enough.
5781                 inner_brackets.add(id(leaf))
5782                 continue
5783
5784             if closing_bracket:
5785                 omit.add(id(closing_bracket))
5786                 omit.update(inner_brackets)
5787                 inner_brackets.clear()
5788                 yield omit
5789
5790             if (
5791                 line.should_explode
5792                 and prev
5793                 and prev.type == token.COMMA
5794                 and not is_one_tuple_between(leaf.opening_bracket, leaf, line.leaves)
5795             ):
5796                 # Never omit bracket pairs with trailing commas.
5797                 # We need to explode on those.
5798                 break
5799
5800             if leaf.value:
5801                 opening_bracket = leaf.opening_bracket
5802                 closing_bracket = leaf
5803
5804
5805 def get_future_imports(node: Node) -> Set[str]:
5806     """Return a set of __future__ imports in the file."""
5807     imports: Set[str] = set()
5808
5809     def get_imports_from_children(children: List[LN]) -> Generator[str, None, None]:
5810         for child in children:
5811             if isinstance(child, Leaf):
5812                 if child.type == token.NAME:
5813                     yield child.value
5814
5815             elif child.type == syms.import_as_name:
5816                 orig_name = child.children[0]
5817                 assert isinstance(orig_name, Leaf), "Invalid syntax parsing imports"
5818                 assert orig_name.type == token.NAME, "Invalid syntax parsing imports"
5819                 yield orig_name.value
5820
5821             elif child.type == syms.import_as_names:
5822                 yield from get_imports_from_children(child.children)
5823
5824             else:
5825                 raise AssertionError("Invalid syntax parsing imports")
5826
5827     for child in node.children:
5828         if child.type != syms.simple_stmt:
5829             break
5830
5831         first_child = child.children[0]
5832         if isinstance(first_child, Leaf):
5833             # Continue looking if we see a docstring; otherwise stop.
5834             if (
5835                 len(child.children) == 2
5836                 and first_child.type == token.STRING
5837                 and child.children[1].type == token.NEWLINE
5838             ):
5839                 continue
5840
5841             break
5842
5843         elif first_child.type == syms.import_from:
5844             module_name = first_child.children[1]
5845             if not isinstance(module_name, Leaf) or module_name.value != "__future__":
5846                 break
5847
5848             imports |= set(get_imports_from_children(first_child.children[3:]))
5849         else:
5850             break
5851
5852     return imports
5853
5854
5855 @lru_cache()
5856 def get_gitignore(root: Path) -> PathSpec:
5857     """ Return a PathSpec matching gitignore content if present."""
5858     gitignore = root / ".gitignore"
5859     lines: List[str] = []
5860     if gitignore.is_file():
5861         with gitignore.open() as gf:
5862             lines = gf.readlines()
5863     return PathSpec.from_lines("gitwildmatch", lines)
5864
5865
5866 def normalize_path_maybe_ignore(
5867     path: Path, root: Path, report: "Report"
5868 ) -> Optional[str]:
5869     """Normalize `path`. May return `None` if `path` was ignored.
5870
5871     `report` is where "path ignored" output goes.
5872     """
5873     try:
5874         abspath = path if path.is_absolute() else Path.cwd() / path
5875         normalized_path = abspath.resolve().relative_to(root).as_posix()
5876     except OSError as e:
5877         report.path_ignored(path, f"cannot be read because {e}")
5878         return None
5879
5880     except ValueError:
5881         if path.is_symlink():
5882             report.path_ignored(path, f"is a symbolic link that points outside {root}")
5883             return None
5884
5885         raise
5886
5887     return normalized_path
5888
5889
5890 def gen_python_files(
5891     paths: Iterable[Path],
5892     root: Path,
5893     include: Optional[Pattern[str]],
5894     exclude: Pattern[str],
5895     force_exclude: Optional[Pattern[str]],
5896     report: "Report",
5897     gitignore: PathSpec,
5898 ) -> Iterator[Path]:
5899     """Generate all files under `path` whose paths are not excluded by the
5900     `exclude_regex` or `force_exclude` regexes, but are included by the `include` regex.
5901
5902     Symbolic links pointing outside of the `root` directory are ignored.
5903
5904     `report` is where output about exclusions goes.
5905     """
5906     assert root.is_absolute(), f"INTERNAL ERROR: `root` must be absolute but is {root}"
5907     for child in paths:
5908         normalized_path = normalize_path_maybe_ignore(child, root, report)
5909         if normalized_path is None:
5910             continue
5911
5912         # First ignore files matching .gitignore
5913         if gitignore.match_file(normalized_path):
5914             report.path_ignored(child, "matches the .gitignore file content")
5915             continue
5916
5917         # Then ignore with `--exclude` and `--force-exclude` options.
5918         normalized_path = "/" + normalized_path
5919         if child.is_dir():
5920             normalized_path += "/"
5921
5922         exclude_match = exclude.search(normalized_path) if exclude else None
5923         if exclude_match and exclude_match.group(0):
5924             report.path_ignored(child, "matches the --exclude regular expression")
5925             continue
5926
5927         force_exclude_match = (
5928             force_exclude.search(normalized_path) if force_exclude else None
5929         )
5930         if force_exclude_match and force_exclude_match.group(0):
5931             report.path_ignored(child, "matches the --force-exclude regular expression")
5932             continue
5933
5934         if child.is_dir():
5935             yield from gen_python_files(
5936                 child.iterdir(),
5937                 root,
5938                 include,
5939                 exclude,
5940                 force_exclude,
5941                 report,
5942                 gitignore,
5943             )
5944
5945         elif child.is_file():
5946             include_match = include.search(normalized_path) if include else True
5947             if include_match:
5948                 yield child
5949
5950
5951 @lru_cache()
5952 def find_project_root(srcs: Iterable[str]) -> Path:
5953     """Return a directory containing .git, .hg, or pyproject.toml.
5954
5955     That directory will be a common parent of all files and directories
5956     passed in `srcs`.
5957
5958     If no directory in the tree contains a marker that would specify it's the
5959     project root, the root of the file system is returned.
5960     """
5961     if not srcs:
5962         return Path("/").resolve()
5963
5964     path_srcs = [Path(Path.cwd(), src).resolve() for src in srcs]
5965
5966     # A list of lists of parents for each 'src'. 'src' is included as a
5967     # "parent" of itself if it is a directory
5968     src_parents = [
5969         list(path.parents) + ([path] if path.is_dir() else []) for path in path_srcs
5970     ]
5971
5972     common_base = max(
5973         set.intersection(*(set(parents) for parents in src_parents)),
5974         key=lambda path: path.parts,
5975     )
5976
5977     for directory in (common_base, *common_base.parents):
5978         if (directory / ".git").exists():
5979             return directory
5980
5981         if (directory / ".hg").is_dir():
5982             return directory
5983
5984         if (directory / "pyproject.toml").is_file():
5985             return directory
5986
5987     return directory
5988
5989
5990 @dataclass
5991 class Report:
5992     """Provides a reformatting counter. Can be rendered with `str(report)`."""
5993
5994     check: bool = False
5995     diff: bool = False
5996     quiet: bool = False
5997     verbose: bool = False
5998     change_count: int = 0
5999     same_count: int = 0
6000     failure_count: int = 0
6001
6002     def done(self, src: Path, changed: Changed) -> None:
6003         """Increment the counter for successful reformatting. Write out a message."""
6004         if changed is Changed.YES:
6005             reformatted = "would reformat" if self.check or self.diff else "reformatted"
6006             if self.verbose or not self.quiet:
6007                 out(f"{reformatted} {src}")
6008             self.change_count += 1
6009         else:
6010             if self.verbose:
6011                 if changed is Changed.NO:
6012                     msg = f"{src} already well formatted, good job."
6013                 else:
6014                     msg = f"{src} wasn't modified on disk since last run."
6015                 out(msg, bold=False)
6016             self.same_count += 1
6017
6018     def failed(self, src: Path, message: str) -> None:
6019         """Increment the counter for failed reformatting. Write out a message."""
6020         err(f"error: cannot format {src}: {message}")
6021         self.failure_count += 1
6022
6023     def path_ignored(self, path: Path, message: str) -> None:
6024         if self.verbose:
6025             out(f"{path} ignored: {message}", bold=False)
6026
6027     @property
6028     def return_code(self) -> int:
6029         """Return the exit code that the app should use.
6030
6031         This considers the current state of changed files and failures:
6032         - if there were any failures, return 123;
6033         - if any files were changed and --check is being used, return 1;
6034         - otherwise return 0.
6035         """
6036         # According to http://tldp.org/LDP/abs/html/exitcodes.html starting with
6037         # 126 we have special return codes reserved by the shell.
6038         if self.failure_count:
6039             return 123
6040
6041         elif self.change_count and self.check:
6042             return 1
6043
6044         return 0
6045
6046     def __str__(self) -> str:
6047         """Render a color report of the current state.
6048
6049         Use `click.unstyle` to remove colors.
6050         """
6051         if self.check or self.diff:
6052             reformatted = "would be reformatted"
6053             unchanged = "would be left unchanged"
6054             failed = "would fail to reformat"
6055         else:
6056             reformatted = "reformatted"
6057             unchanged = "left unchanged"
6058             failed = "failed to reformat"
6059         report = []
6060         if self.change_count:
6061             s = "s" if self.change_count > 1 else ""
6062             report.append(
6063                 click.style(f"{self.change_count} file{s} {reformatted}", bold=True)
6064             )
6065         if self.same_count:
6066             s = "s" if self.same_count > 1 else ""
6067             report.append(f"{self.same_count} file{s} {unchanged}")
6068         if self.failure_count:
6069             s = "s" if self.failure_count > 1 else ""
6070             report.append(
6071                 click.style(f"{self.failure_count} file{s} {failed}", fg="red")
6072             )
6073         return ", ".join(report) + "."
6074
6075
6076 def parse_ast(src: str) -> Union[ast.AST, ast3.AST, ast27.AST]:
6077     filename = "<unknown>"
6078     if sys.version_info >= (3, 8):
6079         # TODO: support Python 4+ ;)
6080         for minor_version in range(sys.version_info[1], 4, -1):
6081             try:
6082                 return ast.parse(src, filename, feature_version=(3, minor_version))
6083             except SyntaxError:
6084                 continue
6085     else:
6086         for feature_version in (7, 6):
6087             try:
6088                 return ast3.parse(src, filename, feature_version=feature_version)
6089             except SyntaxError:
6090                 continue
6091
6092     return ast27.parse(src)
6093
6094
6095 def _fixup_ast_constants(
6096     node: Union[ast.AST, ast3.AST, ast27.AST]
6097 ) -> Union[ast.AST, ast3.AST, ast27.AST]:
6098     """Map ast nodes deprecated in 3.8 to Constant."""
6099     if isinstance(node, (ast.Str, ast3.Str, ast27.Str, ast.Bytes, ast3.Bytes)):
6100         return ast.Constant(value=node.s)
6101
6102     if isinstance(node, (ast.Num, ast3.Num, ast27.Num)):
6103         return ast.Constant(value=node.n)
6104
6105     if isinstance(node, (ast.NameConstant, ast3.NameConstant)):
6106         return ast.Constant(value=node.value)
6107
6108     return node
6109
6110
6111 def _stringify_ast(
6112     node: Union[ast.AST, ast3.AST, ast27.AST], depth: int = 0
6113 ) -> Iterator[str]:
6114     """Simple visitor generating strings to compare ASTs by content."""
6115
6116     node = _fixup_ast_constants(node)
6117
6118     yield f"{'  ' * depth}{node.__class__.__name__}("
6119
6120     for field in sorted(node._fields):  # noqa: F402
6121         # TypeIgnore has only one field 'lineno' which breaks this comparison
6122         type_ignore_classes = (ast3.TypeIgnore, ast27.TypeIgnore)
6123         if sys.version_info >= (3, 8):
6124             type_ignore_classes += (ast.TypeIgnore,)
6125         if isinstance(node, type_ignore_classes):
6126             break
6127
6128         try:
6129             value = getattr(node, field)
6130         except AttributeError:
6131             continue
6132
6133         yield f"{'  ' * (depth+1)}{field}="
6134
6135         if isinstance(value, list):
6136             for item in value:
6137                 # Ignore nested tuples within del statements, because we may insert
6138                 # parentheses and they change the AST.
6139                 if (
6140                     field == "targets"
6141                     and isinstance(node, (ast.Delete, ast3.Delete, ast27.Delete))
6142                     and isinstance(item, (ast.Tuple, ast3.Tuple, ast27.Tuple))
6143                 ):
6144                     for item in item.elts:
6145                         yield from _stringify_ast(item, depth + 2)
6146
6147                 elif isinstance(item, (ast.AST, ast3.AST, ast27.AST)):
6148                     yield from _stringify_ast(item, depth + 2)
6149
6150         elif isinstance(value, (ast.AST, ast3.AST, ast27.AST)):
6151             yield from _stringify_ast(value, depth + 2)
6152
6153         else:
6154             # Constant strings may be indented across newlines, if they are
6155             # docstrings; fold spaces after newlines when comparing. Similarly,
6156             # trailing and leading space may be removed.
6157             if (
6158                 isinstance(node, ast.Constant)
6159                 and field == "value"
6160                 and isinstance(value, str)
6161             ):
6162                 normalized = re.sub(r" *\n[ \t]*", "\n", value).strip()
6163             else:
6164                 normalized = value
6165             yield f"{'  ' * (depth+2)}{normalized!r},  # {value.__class__.__name__}"
6166
6167     yield f"{'  ' * depth})  # /{node.__class__.__name__}"
6168
6169
6170 def assert_equivalent(src: str, dst: str) -> None:
6171     """Raise AssertionError if `src` and `dst` aren't equivalent."""
6172     try:
6173         src_ast = parse_ast(src)
6174     except Exception as exc:
6175         raise AssertionError(
6176             "cannot use --safe with this file; failed to parse source file.  AST"
6177             f" error message: {exc}"
6178         )
6179
6180     try:
6181         dst_ast = parse_ast(dst)
6182     except Exception as exc:
6183         log = dump_to_file("".join(traceback.format_tb(exc.__traceback__)), dst)
6184         raise AssertionError(
6185             f"INTERNAL ERROR: Black produced invalid code: {exc}. Please report a bug"
6186             " on https://github.com/psf/black/issues.  This invalid output might be"
6187             f" helpful: {log}"
6188         ) from None
6189
6190     src_ast_str = "\n".join(_stringify_ast(src_ast))
6191     dst_ast_str = "\n".join(_stringify_ast(dst_ast))
6192     if src_ast_str != dst_ast_str:
6193         log = dump_to_file(diff(src_ast_str, dst_ast_str, "src", "dst"))
6194         raise AssertionError(
6195             "INTERNAL ERROR: Black produced code that is not equivalent to the"
6196             " source.  Please report a bug on https://github.com/psf/black/issues. "
6197             f" This diff might be helpful: {log}"
6198         ) from None
6199
6200
6201 def assert_stable(src: str, dst: str, mode: Mode) -> None:
6202     """Raise AssertionError if `dst` reformats differently the second time."""
6203     newdst = format_str(dst, mode=mode)
6204     if dst != newdst:
6205         log = dump_to_file(
6206             str(mode),
6207             diff(src, dst, "source", "first pass"),
6208             diff(dst, newdst, "first pass", "second pass"),
6209         )
6210         raise AssertionError(
6211             "INTERNAL ERROR: Black produced different code on the second pass of the"
6212             " formatter.  Please report a bug on https://github.com/psf/black/issues."
6213             f"  This diff might be helpful: {log}"
6214         ) from None
6215
6216
6217 @mypyc_attr(patchable=True)
6218 def dump_to_file(*output: str) -> str:
6219     """Dump `output` to a temporary file. Return path to the file."""
6220     with tempfile.NamedTemporaryFile(
6221         mode="w", prefix="blk_", suffix=".log", delete=False, encoding="utf8"
6222     ) as f:
6223         for lines in output:
6224             f.write(lines)
6225             if lines and lines[-1] != "\n":
6226                 f.write("\n")
6227     return f.name
6228
6229
6230 @contextmanager
6231 def nullcontext() -> Iterator[None]:
6232     """Return an empty context manager.
6233
6234     To be used like `nullcontext` in Python 3.7.
6235     """
6236     yield
6237
6238
6239 def diff(a: str, b: str, a_name: str, b_name: str) -> str:
6240     """Return a unified diff string between strings `a` and `b`."""
6241     import difflib
6242
6243     a_lines = [line + "\n" for line in a.splitlines()]
6244     b_lines = [line + "\n" for line in b.splitlines()]
6245     return "".join(
6246         difflib.unified_diff(a_lines, b_lines, fromfile=a_name, tofile=b_name, n=5)
6247     )
6248
6249
6250 def cancel(tasks: Iterable["asyncio.Task[Any]"]) -> None:
6251     """asyncio signal handler that cancels all `tasks` and reports to stderr."""
6252     err("Aborted!")
6253     for task in tasks:
6254         task.cancel()
6255
6256
6257 def shutdown(loop: asyncio.AbstractEventLoop) -> None:
6258     """Cancel all pending tasks on `loop`, wait for them, and close the loop."""
6259     try:
6260         if sys.version_info[:2] >= (3, 7):
6261             all_tasks = asyncio.all_tasks
6262         else:
6263             all_tasks = asyncio.Task.all_tasks
6264         # This part is borrowed from asyncio/runners.py in Python 3.7b2.
6265         to_cancel = [task for task in all_tasks(loop) if not task.done()]
6266         if not to_cancel:
6267             return
6268
6269         for task in to_cancel:
6270             task.cancel()
6271         loop.run_until_complete(
6272             asyncio.gather(*to_cancel, loop=loop, return_exceptions=True)
6273         )
6274     finally:
6275         # `concurrent.futures.Future` objects cannot be cancelled once they
6276         # are already running. There might be some when the `shutdown()` happened.
6277         # Silence their logger's spew about the event loop being closed.
6278         cf_logger = logging.getLogger("concurrent.futures")
6279         cf_logger.setLevel(logging.CRITICAL)
6280         loop.close()
6281
6282
6283 def sub_twice(regex: Pattern[str], replacement: str, original: str) -> str:
6284     """Replace `regex` with `replacement` twice on `original`.
6285
6286     This is used by string normalization to perform replaces on
6287     overlapping matches.
6288     """
6289     return regex.sub(replacement, regex.sub(replacement, original))
6290
6291
6292 def re_compile_maybe_verbose(regex: str) -> Pattern[str]:
6293     """Compile a regular expression string in `regex`.
6294
6295     If it contains newlines, use verbose mode.
6296     """
6297     if "\n" in regex:
6298         regex = "(?x)" + regex
6299     compiled: Pattern[str] = re.compile(regex)
6300     return compiled
6301
6302
6303 def enumerate_reversed(sequence: Sequence[T]) -> Iterator[Tuple[Index, T]]:
6304     """Like `reversed(enumerate(sequence))` if that were possible."""
6305     index = len(sequence) - 1
6306     for element in reversed(sequence):
6307         yield (index, element)
6308         index -= 1
6309
6310
6311 def enumerate_with_length(
6312     line: Line, reversed: bool = False
6313 ) -> Iterator[Tuple[Index, Leaf, int]]:
6314     """Return an enumeration of leaves with their length.
6315
6316     Stops prematurely on multiline strings and standalone comments.
6317     """
6318     op = cast(
6319         Callable[[Sequence[Leaf]], Iterator[Tuple[Index, Leaf]]],
6320         enumerate_reversed if reversed else enumerate,
6321     )
6322     for index, leaf in op(line.leaves):
6323         length = len(leaf.prefix) + len(leaf.value)
6324         if "\n" in leaf.value:
6325             return  # Multiline strings, we can't continue.
6326
6327         for comment in line.comments_after(leaf):
6328             length += len(comment.value)
6329
6330         yield index, leaf, length
6331
6332
6333 def is_line_short_enough(line: Line, *, line_length: int, line_str: str = "") -> bool:
6334     """Return True if `line` is no longer than `line_length`.
6335
6336     Uses the provided `line_str` rendering, if any, otherwise computes a new one.
6337     """
6338     if not line_str:
6339         line_str = line_to_string(line)
6340     return (
6341         len(line_str) <= line_length
6342         and "\n" not in line_str  # multiline strings
6343         and not line.contains_standalone_comments()
6344     )
6345
6346
6347 def can_be_split(line: Line) -> bool:
6348     """Return False if the line cannot be split *for sure*.
6349
6350     This is not an exhaustive search but a cheap heuristic that we can use to
6351     avoid some unfortunate formattings (mostly around wrapping unsplittable code
6352     in unnecessary parentheses).
6353     """
6354     leaves = line.leaves
6355     if len(leaves) < 2:
6356         return False
6357
6358     if leaves[0].type == token.STRING and leaves[1].type == token.DOT:
6359         call_count = 0
6360         dot_count = 0
6361         next = leaves[-1]
6362         for leaf in leaves[-2::-1]:
6363             if leaf.type in OPENING_BRACKETS:
6364                 if next.type not in CLOSING_BRACKETS:
6365                     return False
6366
6367                 call_count += 1
6368             elif leaf.type == token.DOT:
6369                 dot_count += 1
6370             elif leaf.type == token.NAME:
6371                 if not (next.type == token.DOT or next.type in OPENING_BRACKETS):
6372                     return False
6373
6374             elif leaf.type not in CLOSING_BRACKETS:
6375                 return False
6376
6377             if dot_count > 1 and call_count > 1:
6378                 return False
6379
6380     return True
6381
6382
6383 def can_omit_invisible_parens(
6384     line: Line,
6385     line_length: int,
6386     omit_on_explode: Collection[LeafID] = (),
6387 ) -> bool:
6388     """Does `line` have a shape safe to reformat without optional parens around it?
6389
6390     Returns True for only a subset of potentially nice looking formattings but
6391     the point is to not return false positives that end up producing lines that
6392     are too long.
6393     """
6394     bt = line.bracket_tracker
6395     if not bt.delimiters:
6396         # Without delimiters the optional parentheses are useless.
6397         return True
6398
6399     max_priority = bt.max_delimiter_priority()
6400     if bt.delimiter_count_with_priority(max_priority) > 1:
6401         # With more than one delimiter of a kind the optional parentheses read better.
6402         return False
6403
6404     if max_priority == DOT_PRIORITY:
6405         # A single stranded method call doesn't require optional parentheses.
6406         return True
6407
6408     assert len(line.leaves) >= 2, "Stranded delimiter"
6409
6410     # With a single delimiter, omit if the expression starts or ends with
6411     # a bracket.
6412     first = line.leaves[0]
6413     second = line.leaves[1]
6414     if first.type in OPENING_BRACKETS and second.type not in CLOSING_BRACKETS:
6415         if _can_omit_opening_paren(line, first=first, line_length=line_length):
6416             return True
6417
6418         # Note: we are not returning False here because a line might have *both*
6419         # a leading opening bracket and a trailing closing bracket.  If the
6420         # opening bracket doesn't match our rule, maybe the closing will.
6421
6422     penultimate = line.leaves[-2]
6423     last = line.leaves[-1]
6424     if line.should_explode:
6425         try:
6426             penultimate, last = last_two_except(line.leaves, omit=omit_on_explode)
6427         except LookupError:
6428             # Turns out we'd omit everything.  We cannot skip the optional parentheses.
6429             return False
6430
6431     if (
6432         last.type == token.RPAR
6433         or last.type == token.RBRACE
6434         or (
6435             # don't use indexing for omitting optional parentheses;
6436             # it looks weird
6437             last.type == token.RSQB
6438             and last.parent
6439             and last.parent.type != syms.trailer
6440         )
6441     ):
6442         if penultimate.type in OPENING_BRACKETS:
6443             # Empty brackets don't help.
6444             return False
6445
6446         if is_multiline_string(first):
6447             # Additional wrapping of a multiline string in this situation is
6448             # unnecessary.
6449             return True
6450
6451         if line.should_explode and penultimate.type == token.COMMA:
6452             # The rightmost non-omitted bracket pair is the one we want to explode on.
6453             return True
6454
6455         if _can_omit_closing_paren(line, last=last, line_length=line_length):
6456             return True
6457
6458     return False
6459
6460
6461 def _can_omit_opening_paren(line: Line, *, first: Leaf, line_length: int) -> bool:
6462     """See `can_omit_invisible_parens`."""
6463     remainder = False
6464     length = 4 * line.depth
6465     _index = -1
6466     for _index, leaf, leaf_length in enumerate_with_length(line):
6467         if leaf.type in CLOSING_BRACKETS and leaf.opening_bracket is first:
6468             remainder = True
6469         if remainder:
6470             length += leaf_length
6471             if length > line_length:
6472                 break
6473
6474             if leaf.type in OPENING_BRACKETS:
6475                 # There are brackets we can further split on.
6476                 remainder = False
6477
6478     else:
6479         # checked the entire string and line length wasn't exceeded
6480         if len(line.leaves) == _index + 1:
6481             return True
6482
6483     return False
6484
6485
6486 def _can_omit_closing_paren(line: Line, *, last: Leaf, line_length: int) -> bool:
6487     """See `can_omit_invisible_parens`."""
6488     length = 4 * line.depth
6489     seen_other_brackets = False
6490     for _index, leaf, leaf_length in enumerate_with_length(line):
6491         length += leaf_length
6492         if leaf is last.opening_bracket:
6493             if seen_other_brackets or length <= line_length:
6494                 return True
6495
6496         elif leaf.type in OPENING_BRACKETS:
6497             # There are brackets we can further split on.
6498             seen_other_brackets = True
6499
6500     return False
6501
6502
6503 def last_two_except(leaves: List[Leaf], omit: Collection[LeafID]) -> Tuple[Leaf, Leaf]:
6504     """Return (penultimate, last) leaves skipping brackets in `omit` and contents."""
6505     stop_after = None
6506     last = None
6507     for leaf in reversed(leaves):
6508         if stop_after:
6509             if leaf is stop_after:
6510                 stop_after = None
6511             continue
6512
6513         if last:
6514             return leaf, last
6515
6516         if id(leaf) in omit:
6517             stop_after = leaf.opening_bracket
6518         else:
6519             last = leaf
6520     else:
6521         raise LookupError("Last two leaves were also skipped")
6522
6523
6524 def run_transformer(
6525     line: Line,
6526     transform: Transformer,
6527     mode: Mode,
6528     features: Collection[Feature],
6529     *,
6530     line_str: str = "",
6531 ) -> List[Line]:
6532     if not line_str:
6533         line_str = line_to_string(line)
6534     result: List[Line] = []
6535     for transformed_line in transform(line, features):
6536         if str(transformed_line).strip("\n") == line_str:
6537             raise CannotTransform("Line transformer returned an unchanged result")
6538
6539         result.extend(transform_line(transformed_line, mode=mode, features=features))
6540
6541     if not (
6542         transform.__name__ == "rhs"
6543         and line.bracket_tracker.invisible
6544         and not any(bracket.value for bracket in line.bracket_tracker.invisible)
6545         and not line.contains_multiline_strings()
6546         and not result[0].contains_uncollapsable_type_comments()
6547         and not result[0].contains_unsplittable_type_ignore()
6548         and not is_line_short_enough(result[0], line_length=mode.line_length)
6549     ):
6550         return result
6551
6552     line_copy = line.clone()
6553     append_leaves(line_copy, line, line.leaves)
6554     features_fop = set(features) | {Feature.FORCE_OPTIONAL_PARENTHESES}
6555     second_opinion = run_transformer(
6556         line_copy, transform, mode, features_fop, line_str=line_str
6557     )
6558     if all(
6559         is_line_short_enough(ln, line_length=mode.line_length) for ln in second_opinion
6560     ):
6561         result = second_opinion
6562     return result
6563
6564
6565 def get_cache_file(mode: Mode) -> Path:
6566     return CACHE_DIR / f"cache.{mode.get_cache_key()}.pickle"
6567
6568
6569 def read_cache(mode: Mode) -> Cache:
6570     """Read the cache if it exists and is well formed.
6571
6572     If it is not well formed, the call to write_cache later should resolve the issue.
6573     """
6574     cache_file = get_cache_file(mode)
6575     if not cache_file.exists():
6576         return {}
6577
6578     with cache_file.open("rb") as fobj:
6579         try:
6580             cache: Cache = pickle.load(fobj)
6581         except (pickle.UnpicklingError, ValueError):
6582             return {}
6583
6584     return cache
6585
6586
6587 def get_cache_info(path: Path) -> CacheInfo:
6588     """Return the information used to check if a file is already formatted or not."""
6589     stat = path.stat()
6590     return stat.st_mtime, stat.st_size
6591
6592
6593 def filter_cached(cache: Cache, sources: Iterable[Path]) -> Tuple[Set[Path], Set[Path]]:
6594     """Split an iterable of paths in `sources` into two sets.
6595
6596     The first contains paths of files that modified on disk or are not in the
6597     cache. The other contains paths to non-modified files.
6598     """
6599     todo, done = set(), set()
6600     for src in sources:
6601         src = src.resolve()
6602         if cache.get(src) != get_cache_info(src):
6603             todo.add(src)
6604         else:
6605             done.add(src)
6606     return todo, done
6607
6608
6609 def write_cache(cache: Cache, sources: Iterable[Path], mode: Mode) -> None:
6610     """Update the cache file."""
6611     cache_file = get_cache_file(mode)
6612     try:
6613         CACHE_DIR.mkdir(parents=True, exist_ok=True)
6614         new_cache = {**cache, **{src.resolve(): get_cache_info(src) for src in sources}}
6615         with tempfile.NamedTemporaryFile(dir=str(cache_file.parent), delete=False) as f:
6616             pickle.dump(new_cache, f, protocol=4)
6617         os.replace(f.name, cache_file)
6618     except OSError:
6619         pass
6620
6621
6622 def patch_click() -> None:
6623     """Make Click not crash.
6624
6625     On certain misconfigured environments, Python 3 selects the ASCII encoding as the
6626     default which restricts paths that it can access during the lifetime of the
6627     application.  Click refuses to work in this scenario by raising a RuntimeError.
6628
6629     In case of Black the likelihood that non-ASCII characters are going to be used in
6630     file paths is minimal since it's Python source code.  Moreover, this crash was
6631     spurious on Python 3.7 thanks to PEP 538 and PEP 540.
6632     """
6633     try:
6634         from click import core
6635         from click import _unicodefun  # type: ignore
6636     except ModuleNotFoundError:
6637         return
6638
6639     for module in (core, _unicodefun):
6640         if hasattr(module, "_verify_python3_env"):
6641             module._verify_python3_env = lambda: None
6642
6643
6644 def patched_main() -> None:
6645     freeze_support()
6646     patch_click()
6647     main()
6648
6649
6650 def is_docstring(leaf: Leaf) -> bool:
6651     if not is_multiline_string(leaf):
6652         # For the purposes of docstring re-indentation, we don't need to do anything
6653         # with single-line docstrings.
6654         return False
6655
6656     if prev_siblings_are(
6657         leaf.parent, [None, token.NEWLINE, token.INDENT, syms.simple_stmt]
6658     ):
6659         return True
6660
6661     # Multiline docstring on the same line as the `def`.
6662     if prev_siblings_are(leaf.parent, [syms.parameters, token.COLON, syms.simple_stmt]):
6663         # `syms.parameters` is only used in funcdefs and async_funcdefs in the Python
6664         # grammar. We're safe to return True without further checks.
6665         return True
6666
6667     return False
6668
6669
6670 def fix_docstring(docstring: str, prefix: str) -> str:
6671     # https://www.python.org/dev/peps/pep-0257/#handling-docstring-indentation
6672     if not docstring:
6673         return ""
6674     # Convert tabs to spaces (following the normal Python rules)
6675     # and split into a list of lines:
6676     lines = docstring.expandtabs().splitlines()
6677     # Determine minimum indentation (first line doesn't count):
6678     indent = sys.maxsize
6679     for line in lines[1:]:
6680         stripped = line.lstrip()
6681         if stripped:
6682             indent = min(indent, len(line) - len(stripped))
6683     # Remove indentation (first line is special):
6684     trimmed = [lines[0].strip()]
6685     if indent < sys.maxsize:
6686         last_line_idx = len(lines) - 2
6687         for i, line in enumerate(lines[1:]):
6688             stripped_line = line[indent:].rstrip()
6689             if stripped_line or i == last_line_idx:
6690                 trimmed.append(prefix + stripped_line)
6691             else:
6692                 trimmed.append("")
6693     return "\n".join(trimmed)
6694
6695
6696 if __name__ == "__main__":
6697     patched_main()