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

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