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

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