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

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