]> git.madduck.net Git - etc/vim.git/blob - black.py

madduck's git repository

Every one of the projects in this repository is available at the canonical URL git://git.madduck.net/madduck/pub/<projectpath> — see each project's metadata for the exact URL.

All patches and comments are welcome. Please squash your changes to logical commits before using git-format-patch and git-send-email to patches@git.madduck.net. If you'd read over the Git project's submission guidelines and adhered to them, I'd be especially grateful.

SSH access, as well as push access can be individually arranged.

If you use my repositories frequently, consider adding the following snippet to ~/.gitconfig and using the third clone URL listed for each project:

[url "git://git.madduck.net/madduck/"]
  insteadOf = madduck:

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