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

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