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

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