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

madduck's git repository

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

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

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

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

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

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