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

a61a40fe5995e7a41fc0b29e6d3fa052f86f9ace
[etc/vim.git] / black.py
1 #!/usr/bin/env python3
2
3 import asyncio
4 from asyncio.base_events import BaseEventLoop
5 from concurrent.futures import Executor, ProcessPoolExecutor
6 from functools import partial
7 import keyword
8 import os
9 from pathlib import Path
10 import tokenize
11 import sys
12 from typing import (
13     Dict, Generic, Iterable, Iterator, List, Optional, Set, Tuple, Type, TypeVar, Union
14 )
15
16 from attr import dataclass, Factory
17 import click
18
19 # lib2to3 fork
20 from blib2to3.pytree import Node, Leaf, type_repr
21 from blib2to3 import pygram, pytree
22 from blib2to3.pgen2 import driver, token
23 from blib2to3.pgen2.parse import ParseError
24
25 __version__ = "18.3a4"
26 DEFAULT_LINE_LENGTH = 88
27 # types
28 syms = pygram.python_symbols
29 FileContent = str
30 Encoding = str
31 Depth = int
32 NodeType = int
33 LeafID = int
34 Priority = int
35 LN = Union[Leaf, Node]
36 out = partial(click.secho, bold=True, err=True)
37 err = partial(click.secho, fg='red', err=True)
38
39
40 class NothingChanged(UserWarning):
41     """Raised by `format_file` when the reformatted code is the same as source."""
42
43
44 class CannotSplit(Exception):
45     """A readable split that fits the allotted line length is impossible.
46
47     Raised by `left_hand_split()`, `right_hand_split()`, and `delimiter_split()`.
48     """
49
50
51 class FormatError(Exception):
52     """Base exception for `# fmt: on` and `# fmt: off` handling.
53
54     It holds the number of bytes of the prefix consumed before the format
55     control comment appeared.
56     """
57
58     def __init__(self, consumed: int) -> None:
59         super().__init__(consumed)
60         self.consumed = consumed
61
62     def trim_prefix(self, leaf: Leaf) -> None:
63         leaf.prefix = leaf.prefix[self.consumed:]
64
65     def leaf_from_consumed(self, leaf: Leaf) -> Leaf:
66         """Returns a new Leaf from the consumed part of the prefix."""
67         unformatted_prefix = leaf.prefix[:self.consumed]
68         return Leaf(token.NEWLINE, unformatted_prefix)
69
70
71 class FormatOn(FormatError):
72     """Found a comment like `# fmt: on` in the file."""
73
74
75 class FormatOff(FormatError):
76     """Found a comment like `# fmt: off` in the file."""
77
78
79 @click.command()
80 @click.option(
81     '-l',
82     '--line-length',
83     type=int,
84     default=DEFAULT_LINE_LENGTH,
85     help='How many character per line to allow.',
86     show_default=True,
87 )
88 @click.option(
89     '--check',
90     is_flag=True,
91     help=(
92         "Don't write back the files, just return the status.  Return code 0 "
93         "means nothing would change.  Return code 1 means some files would be "
94         "reformatted.  Return code 123 means there was an internal error."
95     ),
96 )
97 @click.option(
98     '--fast/--safe',
99     is_flag=True,
100     help='If --fast given, skip temporary sanity checks. [default: --safe]',
101 )
102 @click.version_option(version=__version__)
103 @click.argument(
104     'src',
105     nargs=-1,
106     type=click.Path(
107         exists=True, file_okay=True, dir_okay=True, readable=True, allow_dash=True
108     ),
109 )
110 @click.pass_context
111 def main(
112     ctx: click.Context, line_length: int, check: bool, fast: bool, src: List[str]
113 ) -> None:
114     """The uncompromising code formatter."""
115     sources: List[Path] = []
116     for s in src:
117         p = Path(s)
118         if p.is_dir():
119             sources.extend(gen_python_files_in_dir(p))
120         elif p.is_file():
121             # if a file was explicitly given, we don't care about its extension
122             sources.append(p)
123         elif s == '-':
124             sources.append(Path('-'))
125         else:
126             err(f'invalid path: {s}')
127     if len(sources) == 0:
128         ctx.exit(0)
129     elif len(sources) == 1:
130         p = sources[0]
131         report = Report(check=check)
132         try:
133             if not p.is_file() and str(p) == '-':
134                 changed = format_stdin_to_stdout(
135                     line_length=line_length, fast=fast, write_back=not check
136                 )
137             else:
138                 changed = format_file_in_place(
139                     p, line_length=line_length, fast=fast, write_back=not check
140                 )
141             report.done(p, changed)
142         except Exception as exc:
143             report.failed(p, str(exc))
144         ctx.exit(report.return_code)
145     else:
146         loop = asyncio.get_event_loop()
147         executor = ProcessPoolExecutor(max_workers=os.cpu_count())
148         return_code = 1
149         try:
150             return_code = loop.run_until_complete(
151                 schedule_formatting(
152                     sources, line_length, not check, fast, loop, executor
153                 )
154             )
155         finally:
156             loop.close()
157             ctx.exit(return_code)
158
159
160 async def schedule_formatting(
161     sources: List[Path],
162     line_length: int,
163     write_back: bool,
164     fast: bool,
165     loop: BaseEventLoop,
166     executor: Executor,
167 ) -> int:
168     tasks = {
169         src: loop.run_in_executor(
170             executor, format_file_in_place, src, line_length, fast, write_back
171         )
172         for src in sources
173     }
174     await asyncio.wait(tasks.values())
175     cancelled = []
176     report = Report()
177     for src, task in tasks.items():
178         if not task.done():
179             report.failed(src, 'timed out, cancelling')
180             task.cancel()
181             cancelled.append(task)
182         elif task.exception():
183             report.failed(src, str(task.exception()))
184         else:
185             report.done(src, task.result())
186     if cancelled:
187         await asyncio.wait(cancelled, timeout=2)
188     out('All done! ✨ 🍰 ✨')
189     click.echo(str(report))
190     return report.return_code
191
192
193 def format_file_in_place(
194     src: Path, line_length: int, fast: bool, write_back: bool = False
195 ) -> bool:
196     """Format the file and rewrite if changed. Return True if changed."""
197     with tokenize.open(src) as src_buffer:
198         src_contents = src_buffer.read()
199     try:
200         contents = format_file_contents(
201             src_contents, line_length=line_length, fast=fast
202         )
203     except NothingChanged:
204         return False
205
206     if write_back:
207         with open(src, "w", encoding=src_buffer.encoding) as f:
208             f.write(contents)
209     return True
210
211
212 def format_stdin_to_stdout(
213     line_length: int, fast: bool, write_back: bool = False
214 ) -> bool:
215     """Format file on stdin and pipe output to stdout. Return True if changed."""
216     contents = sys.stdin.read()
217     try:
218         contents = format_file_contents(contents, line_length=line_length, fast=fast)
219         return True
220
221     except NothingChanged:
222         return False
223
224     finally:
225         if write_back:
226             sys.stdout.write(contents)
227
228
229 def format_file_contents(
230     src_contents: str, line_length: int, fast: bool
231 ) -> FileContent:
232     """Reformats a file and returns its contents and encoding."""
233     if src_contents.strip() == '':
234         raise NothingChanged
235
236     dst_contents = format_str(src_contents, line_length=line_length)
237     if src_contents == dst_contents:
238         raise NothingChanged
239
240     if not fast:
241         assert_equivalent(src_contents, dst_contents)
242         assert_stable(src_contents, dst_contents, line_length=line_length)
243     return dst_contents
244
245
246 def format_str(src_contents: str, line_length: int) -> FileContent:
247     """Reformats a string and returns new contents."""
248     src_node = lib2to3_parse(src_contents)
249     dst_contents = ""
250     lines = LineGenerator()
251     elt = EmptyLineTracker()
252     py36 = is_python36(src_node)
253     empty_line = Line()
254     after = 0
255     for current_line in lines.visit(src_node):
256         for _ in range(after):
257             dst_contents += str(empty_line)
258         before, after = elt.maybe_empty_lines(current_line)
259         for _ in range(before):
260             dst_contents += str(empty_line)
261         for line in split_line(current_line, line_length=line_length, py36=py36):
262             dst_contents += str(line)
263     return dst_contents
264
265
266 GRAMMARS = [
267     pygram.python_grammar_no_print_statement_no_exec_statement,
268     pygram.python_grammar_no_print_statement,
269     pygram.python_grammar_no_exec_statement,
270     pygram.python_grammar,
271 ]
272
273
274 def lib2to3_parse(src_txt: str) -> Node:
275     """Given a string with source, return the lib2to3 Node."""
276     grammar = pygram.python_grammar_no_print_statement
277     if src_txt[-1] != '\n':
278         nl = '\r\n' if '\r\n' in src_txt[:1024] else '\n'
279         src_txt += nl
280     for grammar in GRAMMARS:
281         drv = driver.Driver(grammar, pytree.convert)
282         try:
283             result = drv.parse_string(src_txt, True)
284             break
285
286         except ParseError as pe:
287             lineno, column = pe.context[1]
288             lines = src_txt.splitlines()
289             try:
290                 faulty_line = lines[lineno - 1]
291             except IndexError:
292                 faulty_line = "<line number missing in source>"
293             exc = ValueError(f"Cannot parse: {lineno}:{column}: {faulty_line}")
294     else:
295         raise exc from None
296
297     if isinstance(result, Leaf):
298         result = Node(syms.file_input, [result])
299     return result
300
301
302 def lib2to3_unparse(node: Node) -> str:
303     """Given a lib2to3 node, return its string representation."""
304     code = str(node)
305     return code
306
307
308 T = TypeVar('T')
309
310
311 class Visitor(Generic[T]):
312     """Basic lib2to3 visitor that yields things of type `T` on `visit()`."""
313
314     def visit(self, node: LN) -> Iterator[T]:
315         """Main method to start the visit process. Yields objects of type `T`.
316
317         It tries to find a `visit_*()` method for the given `node.type`, like
318         `visit_simple_stmt` for Node objects or `visit_INDENT` for Leaf objects.
319         If no dedicated `visit_*()` method is found, chooses `visit_default()`
320         instead.
321
322         Then yields objects of type `T` from the selected visitor.
323         """
324         if node.type < 256:
325             name = token.tok_name[node.type]
326         else:
327             name = type_repr(node.type)
328         yield from getattr(self, f'visit_{name}', self.visit_default)(node)
329
330     def visit_default(self, node: LN) -> Iterator[T]:
331         """Default `visit_*()` implementation. Recurses to children of `node`."""
332         if isinstance(node, Node):
333             for child in node.children:
334                 yield from self.visit(child)
335
336
337 @dataclass
338 class DebugVisitor(Visitor[T]):
339     tree_depth: int = 0
340
341     def visit_default(self, node: LN) -> Iterator[T]:
342         indent = ' ' * (2 * self.tree_depth)
343         if isinstance(node, Node):
344             _type = type_repr(node.type)
345             out(f'{indent}{_type}', fg='yellow')
346             self.tree_depth += 1
347             for child in node.children:
348                 yield from self.visit(child)
349
350             self.tree_depth -= 1
351             out(f'{indent}/{_type}', fg='yellow', bold=False)
352         else:
353             _type = token.tok_name.get(node.type, str(node.type))
354             out(f'{indent}{_type}', fg='blue', nl=False)
355             if node.prefix:
356                 # We don't have to handle prefixes for `Node` objects since
357                 # that delegates to the first child anyway.
358                 out(f' {node.prefix!r}', fg='green', bold=False, nl=False)
359             out(f' {node.value!r}', fg='blue', bold=False)
360
361     @classmethod
362     def show(cls, code: str) -> None:
363         """Pretty-prints a given string of `code`.
364
365         Convenience method for debugging.
366         """
367         v: DebugVisitor[None] = DebugVisitor()
368         list(v.visit(lib2to3_parse(code)))
369
370
371 KEYWORDS = set(keyword.kwlist)
372 WHITESPACE = {token.DEDENT, token.INDENT, token.NEWLINE}
373 FLOW_CONTROL = {'return', 'raise', 'break', 'continue'}
374 STATEMENT = {
375     syms.if_stmt,
376     syms.while_stmt,
377     syms.for_stmt,
378     syms.try_stmt,
379     syms.except_clause,
380     syms.with_stmt,
381     syms.funcdef,
382     syms.classdef,
383 }
384 STANDALONE_COMMENT = 153
385 LOGIC_OPERATORS = {'and', 'or'}
386 COMPARATORS = {
387     token.LESS,
388     token.GREATER,
389     token.EQEQUAL,
390     token.NOTEQUAL,
391     token.LESSEQUAL,
392     token.GREATEREQUAL,
393 }
394 MATH_OPERATORS = {
395     token.PLUS,
396     token.MINUS,
397     token.STAR,
398     token.SLASH,
399     token.VBAR,
400     token.AMPER,
401     token.PERCENT,
402     token.CIRCUMFLEX,
403     token.TILDE,
404     token.LEFTSHIFT,
405     token.RIGHTSHIFT,
406     token.DOUBLESTAR,
407     token.DOUBLESLASH,
408 }
409 COMPREHENSION_PRIORITY = 20
410 COMMA_PRIORITY = 10
411 LOGIC_PRIORITY = 5
412 STRING_PRIORITY = 4
413 COMPARATOR_PRIORITY = 3
414 MATH_PRIORITY = 1
415
416
417 @dataclass
418 class BracketTracker:
419     """Keeps track of brackets on a line."""
420
421     #: Current bracket depth.
422     depth: int = 0
423     #: All currently unclosed brackets.
424     bracket_match: Dict[Tuple[Depth, NodeType], Leaf] = Factory(dict)
425     #: All current delimiters with their assigned priority.
426     delimiters: Dict[LeafID, Priority] = Factory(dict)
427     #: Last processed leaf, if any.
428     previous: Optional[Leaf] = None
429
430     def mark(self, leaf: Leaf) -> None:
431         """Marks `leaf` with bracket-related metadata. Keeps track of delimiters.
432
433         All leaves receive an int `bracket_depth` field that stores how deep
434         within brackets a given leaf is. 0 means there are no enclosing brackets
435         that started on this line.
436
437         If a leaf is itself a closing bracket, it receives an `opening_bracket`
438         field that it forms a pair with. This is a one-directional link to
439         avoid reference cycles.
440
441         If a leaf is a delimiter (a token on which Black can split the line if
442         needed) and it's on depth 0, its `id()` is stored in the tracker's
443         `delimiters` field.
444         """
445         if leaf.type == token.COMMENT:
446             return
447
448         if leaf.type in CLOSING_BRACKETS:
449             self.depth -= 1
450             opening_bracket = self.bracket_match.pop((self.depth, leaf.type))
451             leaf.opening_bracket = opening_bracket
452         leaf.bracket_depth = self.depth
453         if self.depth == 0:
454             delim = is_delimiter(leaf)
455             if delim:
456                 self.delimiters[id(leaf)] = delim
457             elif self.previous is not None:
458                 if leaf.type == token.STRING and self.previous.type == token.STRING:
459                     self.delimiters[id(self.previous)] = STRING_PRIORITY
460                 elif (
461                     leaf.type == token.NAME
462                     and leaf.value == 'for'
463                     and leaf.parent
464                     and leaf.parent.type in {syms.comp_for, syms.old_comp_for}
465                 ):
466                     self.delimiters[id(self.previous)] = COMPREHENSION_PRIORITY
467                 elif (
468                     leaf.type == token.NAME
469                     and leaf.value == 'if'
470                     and leaf.parent
471                     and leaf.parent.type in {syms.comp_if, syms.old_comp_if}
472                 ):
473                     self.delimiters[id(self.previous)] = COMPREHENSION_PRIORITY
474                 elif (
475                     leaf.type == token.NAME
476                     and leaf.value in LOGIC_OPERATORS
477                     and leaf.parent
478                 ):
479                     self.delimiters[id(self.previous)] = LOGIC_PRIORITY
480         if leaf.type in OPENING_BRACKETS:
481             self.bracket_match[self.depth, BRACKET[leaf.type]] = leaf
482             self.depth += 1
483         self.previous = leaf
484
485     def any_open_brackets(self) -> bool:
486         """Returns True if there is an yet unmatched open bracket on the line."""
487         return bool(self.bracket_match)
488
489     def max_delimiter_priority(self, exclude: Iterable[LeafID] = ()) -> int:
490         """Returns the highest priority of a delimiter found on the line.
491
492         Values are consistent with what `is_delimiter()` returns.
493         """
494         return max(v for k, v in self.delimiters.items() if k not in exclude)
495
496
497 @dataclass
498 class Line:
499     """Holds leaves and comments. Can be printed with `str(line)`."""
500
501     #: indentation level
502     depth: int = 0
503     #: list of leaves
504     leaves: List[Leaf] = Factory(list)
505     #: inline comments that belong on this line
506     comments: Dict[LeafID, Leaf] = Factory(dict)
507     bracket_tracker: BracketTracker = Factory(BracketTracker)
508     inside_brackets: bool = False
509     has_for: bool = False
510     _for_loop_variable: bool = False
511
512     def append(self, leaf: Leaf, preformatted: bool = False) -> None:
513         """Add a new `leaf` to the end of the line.
514
515         Unless `preformatted` is True, the `leaf` will receive a new consistent
516         whitespace prefix and metadata applied by :class:`BracketTracker`.
517         Trailing commas are maybe removed, unpacked for loop variables are
518         demoted from being delimiters.
519
520         Inline comments are put aside.
521         """
522         has_value = leaf.value.strip()
523         if not has_value:
524             return
525
526         if self.leaves and not preformatted:
527             # Note: at this point leaf.prefix should be empty except for
528             # imports, for which we only preserve newlines.
529             leaf.prefix += whitespace(leaf)
530         if self.inside_brackets or not preformatted:
531             self.maybe_decrement_after_for_loop_variable(leaf)
532             self.bracket_tracker.mark(leaf)
533             self.maybe_remove_trailing_comma(leaf)
534             self.maybe_increment_for_loop_variable(leaf)
535             if self.maybe_adapt_standalone_comment(leaf):
536                 return
537
538         if not self.append_comment(leaf):
539             self.leaves.append(leaf)
540
541     @property
542     def is_comment(self) -> bool:
543         """Is this line a standalone comment?"""
544         return bool(self) and self.leaves[0].type == STANDALONE_COMMENT
545
546     @property
547     def is_decorator(self) -> bool:
548         """Is this line a decorator?"""
549         return bool(self) and self.leaves[0].type == token.AT
550
551     @property
552     def is_import(self) -> bool:
553         """Is this an import line?"""
554         return bool(self) and is_import(self.leaves[0])
555
556     @property
557     def is_class(self) -> bool:
558         """Is this a class definition?"""
559         return (
560             bool(self)
561             and self.leaves[0].type == token.NAME
562             and self.leaves[0].value == 'class'
563         )
564
565     @property
566     def is_def(self) -> bool:
567         """Is this a function definition? (Also returns True for async defs.)"""
568         try:
569             first_leaf = self.leaves[0]
570         except IndexError:
571             return False
572
573         try:
574             second_leaf: Optional[Leaf] = self.leaves[1]
575         except IndexError:
576             second_leaf = None
577         return (
578             (first_leaf.type == token.NAME and first_leaf.value == 'def')
579             or (
580                 first_leaf.type == token.ASYNC
581                 and second_leaf is not None
582                 and second_leaf.type == token.NAME
583                 and second_leaf.value == 'def'
584             )
585         )
586
587     @property
588     def is_flow_control(self) -> bool:
589         """Is this a flow control statement?
590
591         Those are `return`, `raise`, `break`, and `continue`.
592         """
593         return (
594             bool(self)
595             and self.leaves[0].type == token.NAME
596             and self.leaves[0].value in FLOW_CONTROL
597         )
598
599     @property
600     def is_yield(self) -> bool:
601         """Is this a yield statement?"""
602         return (
603             bool(self)
604             and self.leaves[0].type == token.NAME
605             and self.leaves[0].value == 'yield'
606         )
607
608     def maybe_remove_trailing_comma(self, closing: Leaf) -> bool:
609         """Remove trailing comma if there is one and it's safe."""
610         if not (
611             self.leaves
612             and self.leaves[-1].type == token.COMMA
613             and closing.type in CLOSING_BRACKETS
614         ):
615             return False
616
617         if closing.type == token.RBRACE:
618             self.leaves.pop()
619             return True
620
621         if closing.type == token.RSQB:
622             comma = self.leaves[-1]
623             if comma.parent and comma.parent.type == syms.listmaker:
624                 self.leaves.pop()
625                 return True
626
627         # For parens let's check if it's safe to remove the comma.  If the
628         # trailing one is the only one, we might mistakenly change a tuple
629         # into a different type by removing the comma.
630         depth = closing.bracket_depth + 1
631         commas = 0
632         opening = closing.opening_bracket
633         for _opening_index, leaf in enumerate(self.leaves):
634             if leaf is opening:
635                 break
636
637         else:
638             return False
639
640         for leaf in self.leaves[_opening_index + 1:]:
641             if leaf is closing:
642                 break
643
644             bracket_depth = leaf.bracket_depth
645             if bracket_depth == depth and leaf.type == token.COMMA:
646                 commas += 1
647                 if leaf.parent and leaf.parent.type == syms.arglist:
648                     commas += 1
649                     break
650
651         if commas > 1:
652             self.leaves.pop()
653             return True
654
655         return False
656
657     def maybe_increment_for_loop_variable(self, leaf: Leaf) -> bool:
658         """In a for loop, or comprehension, the variables are often unpacks.
659
660         To avoid splitting on the comma in this situation, we will increase
661         the depth of tokens between `for` and `in`.
662         """
663         if leaf.type == token.NAME and leaf.value == 'for':
664             self.has_for = True
665             self.bracket_tracker.depth += 1
666             self._for_loop_variable = True
667             return True
668
669         return False
670
671     def maybe_decrement_after_for_loop_variable(self, leaf: Leaf) -> bool:
672         """See `maybe_increment_for_loop_variable` above for explanation."""
673         if self._for_loop_variable and leaf.type == token.NAME and leaf.value == 'in':
674             self.bracket_tracker.depth -= 1
675             self._for_loop_variable = False
676             return True
677
678         return False
679
680     def maybe_adapt_standalone_comment(self, comment: Leaf) -> bool:
681         """Hack a standalone comment to act as a trailing comment for line splitting.
682
683         If this line has brackets and a standalone `comment`, we need to adapt
684         it to be able to still reformat the line.
685
686         This is not perfect, the line to which the standalone comment gets
687         appended will appear "too long" when splitting.
688         """
689         if not (
690             comment.type == STANDALONE_COMMENT
691             and self.bracket_tracker.any_open_brackets()
692         ):
693             return False
694
695         comment.type = token.COMMENT
696         comment.prefix = '\n' + '    ' * (self.depth + 1)
697         return self.append_comment(comment)
698
699     def append_comment(self, comment: Leaf) -> bool:
700         """Add an inline comment to the line."""
701         if comment.type != token.COMMENT:
702             return False
703
704         try:
705             after = id(self.last_non_delimiter())
706         except LookupError:
707             comment.type = STANDALONE_COMMENT
708             comment.prefix = ''
709             return False
710
711         else:
712             if after in self.comments:
713                 self.comments[after].value += str(comment)
714             else:
715                 self.comments[after] = comment
716             return True
717
718     def last_non_delimiter(self) -> Leaf:
719         """Returns the last non-delimiter on the line. Raises LookupError otherwise."""
720         for i in range(len(self.leaves)):
721             last = self.leaves[-i - 1]
722             if not is_delimiter(last):
723                 return last
724
725         raise LookupError("No non-delimiters found")
726
727     def __str__(self) -> str:
728         """Render the line."""
729         if not self:
730             return '\n'
731
732         indent = '    ' * self.depth
733         leaves = iter(self.leaves)
734         first = next(leaves)
735         res = f'{first.prefix}{indent}{first.value}'
736         for leaf in leaves:
737             res += str(leaf)
738         for comment in self.comments.values():
739             res += str(comment)
740         return res + '\n'
741
742     def __bool__(self) -> bool:
743         """Returns True if the line has leaves or comments."""
744         return bool(self.leaves or self.comments)
745
746
747 class UnformattedLines(Line):
748     """Just like :class:`Line` but stores lines which aren't reformatted."""
749
750     def append(self, leaf: Leaf, preformatted: bool = True) -> None:
751         """Just add a new `leaf` to the end of the lines.
752
753         The `preformatted` argument is ignored.
754
755         Keeps track of indentation `depth`, which is useful when the user
756         says `# fmt: on`. Otherwise, doesn't do anything with the `leaf`.
757         """
758         try:
759             list(generate_comments(leaf))
760         except FormatOn as f_on:
761             self.leaves.append(f_on.leaf_from_consumed(leaf))
762             raise
763
764         self.leaves.append(leaf)
765         if leaf.type == token.INDENT:
766             self.depth += 1
767         elif leaf.type == token.DEDENT:
768             self.depth -= 1
769
770     def append_comment(self, comment: Leaf) -> bool:
771         """Not implemented in this class."""
772         raise NotImplementedError("Unformatted lines don't store comments separately.")
773
774     def maybe_remove_trailing_comma(self, closing: Leaf) -> bool:
775         """Does nothing and returns False."""
776         return False
777
778     def maybe_increment_for_loop_variable(self, leaf: Leaf) -> bool:
779         """Does nothing and returns False."""
780         return False
781
782     def maybe_adapt_standalone_comment(self, comment: Leaf) -> bool:
783         """Does nothing and returns False."""
784         return False
785
786     def __str__(self) -> str:
787         """Renders unformatted lines from leaves which were added with `append()`.
788
789         `depth` is not used for indentation in this case.
790         """
791         if not self:
792             return '\n'
793
794         res = ''
795         for leaf in self.leaves:
796             res += str(leaf)
797         return res
798
799
800 @dataclass
801 class EmptyLineTracker:
802     """Provides a stateful method that returns the number of potential extra
803     empty lines needed before and after the currently processed line.
804
805     Note: this tracker works on lines that haven't been split yet.  It assumes
806     the prefix of the first leaf consists of optional newlines.  Those newlines
807     are consumed by `maybe_empty_lines()` and included in the computation.
808     """
809     previous_line: Optional[Line] = None
810     previous_after: int = 0
811     previous_defs: List[int] = Factory(list)
812
813     def maybe_empty_lines(self, current_line: Line) -> Tuple[int, int]:
814         """Returns the number of extra empty lines before and after the `current_line`.
815
816         This is for separating `def`, `async def` and `class` with extra empty lines
817         (two on module-level), as well as providing an extra empty line after flow
818         control keywords to make them more prominent.
819         """
820         if isinstance(current_line, UnformattedLines):
821             return 0, 0
822
823         before, after = self._maybe_empty_lines(current_line)
824         before -= self.previous_after
825         self.previous_after = after
826         self.previous_line = current_line
827         return before, after
828
829     def _maybe_empty_lines(self, current_line: Line) -> Tuple[int, int]:
830         max_allowed = 1
831         if current_line.depth == 0:
832             max_allowed = 2
833         if current_line.leaves:
834             # Consume the first leaf's extra newlines.
835             first_leaf = current_line.leaves[0]
836             before = first_leaf.prefix.count('\n')
837             before = min(before, max_allowed)
838             first_leaf.prefix = ''
839         else:
840             before = 0
841         depth = current_line.depth
842         while self.previous_defs and self.previous_defs[-1] >= depth:
843             self.previous_defs.pop()
844             before = 1 if depth else 2
845         is_decorator = current_line.is_decorator
846         if is_decorator or current_line.is_def or current_line.is_class:
847             if not is_decorator:
848                 self.previous_defs.append(depth)
849             if self.previous_line is None:
850                 # Don't insert empty lines before the first line in the file.
851                 return 0, 0
852
853             if self.previous_line and self.previous_line.is_decorator:
854                 # Don't insert empty lines between decorators.
855                 return 0, 0
856
857             newlines = 2
858             if current_line.depth:
859                 newlines -= 1
860             return newlines, 0
861
862         if current_line.is_flow_control:
863             return before, 1
864
865         if (
866             self.previous_line
867             and self.previous_line.is_import
868             and not current_line.is_import
869             and depth == self.previous_line.depth
870         ):
871             return (before or 1), 0
872
873         if (
874             self.previous_line
875             and self.previous_line.is_yield
876             and (not current_line.is_yield or depth != self.previous_line.depth)
877         ):
878             return (before or 1), 0
879
880         return before, 0
881
882
883 @dataclass
884 class LineGenerator(Visitor[Line]):
885     """Generates reformatted Line objects.  Empty lines are not emitted.
886
887     Note: destroys the tree it's visiting by mutating prefixes of its leaves
888     in ways that will no longer stringify to valid Python code on the tree.
889     """
890     current_line: Line = Factory(Line)
891
892     def line(self, indent: int = 0, type: Type[Line] = Line) -> Iterator[Line]:
893         """Generate a line.
894
895         If the line is empty, only emit if it makes sense.
896         If the line is too long, split it first and then generate.
897
898         If any lines were generated, set up a new current_line.
899         """
900         if not self.current_line:
901             if self.current_line.__class__ == type:
902                 self.current_line.depth += indent
903             else:
904                 self.current_line = type(depth=self.current_line.depth + indent)
905             return  # Line is empty, don't emit. Creating a new one unnecessary.
906
907         complete_line = self.current_line
908         self.current_line = type(depth=complete_line.depth + indent)
909         yield complete_line
910
911     def visit(self, node: LN) -> Iterator[Line]:
912         """Main method to start the visit process. Yields :class:`Line` objects."""
913         if isinstance(self.current_line, UnformattedLines):
914             # File contained `# fmt: off`
915             yield from self.visit_unformatted(node)
916
917         else:
918             yield from super().visit(node)
919
920     def visit_default(self, node: LN) -> Iterator[Line]:
921         """Default `visit_*()` implementation. Recurses to children of `node`."""
922         if isinstance(node, Leaf):
923             any_open_brackets = self.current_line.bracket_tracker.any_open_brackets()
924             try:
925                 for comment in generate_comments(node):
926                     if any_open_brackets:
927                         # any comment within brackets is subject to splitting
928                         self.current_line.append(comment)
929                     elif comment.type == token.COMMENT:
930                         # regular trailing comment
931                         self.current_line.append(comment)
932                         yield from self.line()
933
934                     else:
935                         # regular standalone comment
936                         yield from self.line()
937
938                         self.current_line.append(comment)
939                         yield from self.line()
940
941             except FormatOff as f_off:
942                 f_off.trim_prefix(node)
943                 yield from self.line(type=UnformattedLines)
944                 yield from self.visit(node)
945
946             except FormatOn as f_on:
947                 # This only happens here if somebody says "fmt: on" multiple
948                 # times in a row.
949                 f_on.trim_prefix(node)
950                 yield from self.visit_default(node)
951
952             else:
953                 normalize_prefix(node, inside_brackets=any_open_brackets)
954                 if node.type not in WHITESPACE:
955                     self.current_line.append(node)
956         yield from super().visit_default(node)
957
958     def visit_INDENT(self, node: Node) -> Iterator[Line]:
959         """Increases indentation level, maybe yields a line."""
960         # In blib2to3 INDENT never holds comments.
961         yield from self.line(+1)
962         yield from self.visit_default(node)
963
964     def visit_DEDENT(self, node: Node) -> Iterator[Line]:
965         """Decreases indentation level, maybe yields a line."""
966         # DEDENT has no value. Additionally, in blib2to3 it never holds comments.
967         yield from self.line(-1)
968
969     def visit_stmt(self, node: Node, keywords: Set[str]) -> Iterator[Line]:
970         """Visits a statement.
971
972         This implementation is shared for `if`, `while`, `for`, `try`, `except`,
973         `def`, `with`, and `class`.
974
975         The relevant Python language `keywords` for a given statement will be NAME
976         leaves within it. This methods puts those on a separate line.
977         """
978         for child in node.children:
979             if child.type == token.NAME and child.value in keywords:  # type: ignore
980                 yield from self.line()
981
982             yield from self.visit(child)
983
984     def visit_simple_stmt(self, node: Node) -> Iterator[Line]:
985         """Visits a statement without nested statements."""
986         is_suite_like = node.parent and node.parent.type in STATEMENT
987         if is_suite_like:
988             yield from self.line(+1)
989             yield from self.visit_default(node)
990             yield from self.line(-1)
991
992         else:
993             yield from self.line()
994             yield from self.visit_default(node)
995
996     def visit_async_stmt(self, node: Node) -> Iterator[Line]:
997         """Visits `async def`, `async for`, `async with`."""
998         yield from self.line()
999
1000         children = iter(node.children)
1001         for child in children:
1002             yield from self.visit(child)
1003
1004             if child.type == token.ASYNC:
1005                 break
1006
1007         internal_stmt = next(children)
1008         for child in internal_stmt.children:
1009             yield from self.visit(child)
1010
1011     def visit_decorators(self, node: Node) -> Iterator[Line]:
1012         """Visits decorators."""
1013         for child in node.children:
1014             yield from self.line()
1015             yield from self.visit(child)
1016
1017     def visit_SEMI(self, leaf: Leaf) -> Iterator[Line]:
1018         """Semicolons are always removed.
1019
1020         Statements between them are put on separate lines.
1021         """
1022         yield from self.line()
1023
1024     def visit_ENDMARKER(self, leaf: Leaf) -> Iterator[Line]:
1025         """End of file.
1026
1027         Process outstanding comments and end with a newline.
1028         """
1029         yield from self.visit_default(leaf)
1030         yield from self.line()
1031
1032     def visit_unformatted(self, node: LN) -> Iterator[Line]:
1033         """Used when file contained a `# fmt: off`."""
1034         if isinstance(node, Node):
1035             for child in node.children:
1036                 yield from self.visit(child)
1037
1038         else:
1039             try:
1040                 self.current_line.append(node)
1041             except FormatOn as f_on:
1042                 f_on.trim_prefix(node)
1043                 yield from self.line()
1044                 yield from self.visit(node)
1045
1046     def __attrs_post_init__(self) -> None:
1047         """You are in a twisty little maze of passages."""
1048         v = self.visit_stmt
1049         self.visit_if_stmt = partial(v, keywords={'if', 'else', 'elif'})
1050         self.visit_while_stmt = partial(v, keywords={'while', 'else'})
1051         self.visit_for_stmt = partial(v, keywords={'for', 'else'})
1052         self.visit_try_stmt = partial(v, keywords={'try', 'except', 'else', 'finally'})
1053         self.visit_except_clause = partial(v, keywords={'except'})
1054         self.visit_funcdef = partial(v, keywords={'def'})
1055         self.visit_with_stmt = partial(v, keywords={'with'})
1056         self.visit_classdef = partial(v, keywords={'class'})
1057         self.visit_async_funcdef = self.visit_async_stmt
1058         self.visit_decorated = self.visit_decorators
1059
1060
1061 BRACKET = {token.LPAR: token.RPAR, token.LSQB: token.RSQB, token.LBRACE: token.RBRACE}
1062 OPENING_BRACKETS = set(BRACKET.keys())
1063 CLOSING_BRACKETS = set(BRACKET.values())
1064 BRACKETS = OPENING_BRACKETS | CLOSING_BRACKETS
1065 ALWAYS_NO_SPACE = CLOSING_BRACKETS | {token.COMMA, STANDALONE_COMMENT}
1066
1067
1068 def whitespace(leaf: Leaf) -> str:  # noqa C901
1069     """Return whitespace prefix if needed for the given `leaf`."""
1070     NO = ''
1071     SPACE = ' '
1072     DOUBLESPACE = '  '
1073     t = leaf.type
1074     p = leaf.parent
1075     v = leaf.value
1076     if t in ALWAYS_NO_SPACE:
1077         return NO
1078
1079     if t == token.COMMENT:
1080         return DOUBLESPACE
1081
1082     assert p is not None, f"INTERNAL ERROR: hand-made leaf without parent: {leaf!r}"
1083     if t == token.COLON and p.type not in {syms.subscript, syms.subscriptlist}:
1084         return NO
1085
1086     prev = leaf.prev_sibling
1087     if not prev:
1088         prevp = preceding_leaf(p)
1089         if not prevp or prevp.type in OPENING_BRACKETS:
1090             return NO
1091
1092         if t == token.COLON:
1093             return SPACE if prevp.type == token.COMMA else NO
1094
1095         if prevp.type == token.EQUAL:
1096             if prevp.parent:
1097                 if prevp.parent.type in {
1098                     syms.arglist, syms.argument, syms.parameters, syms.varargslist
1099                 }:
1100                     return NO
1101
1102                 elif prevp.parent.type == syms.typedargslist:
1103                     # A bit hacky: if the equal sign has whitespace, it means we
1104                     # previously found it's a typed argument.  So, we're using
1105                     # that, too.
1106                     return prevp.prefix
1107
1108         elif prevp.type == token.DOUBLESTAR:
1109             if prevp.parent and prevp.parent.type in {
1110                 syms.arglist,
1111                 syms.argument,
1112                 syms.dictsetmaker,
1113                 syms.parameters,
1114                 syms.typedargslist,
1115                 syms.varargslist,
1116             }:
1117                 return NO
1118
1119         elif prevp.type == token.COLON:
1120             if prevp.parent and prevp.parent.type in {syms.subscript, syms.sliceop}:
1121                 return NO
1122
1123         elif (
1124             prevp.parent
1125             and prevp.parent.type in {syms.factor, syms.star_expr}
1126             and prevp.type in MATH_OPERATORS
1127         ):
1128             return NO
1129
1130         elif (
1131             prevp.type == token.RIGHTSHIFT
1132             and prevp.parent
1133             and prevp.parent.type == syms.shift_expr
1134             and prevp.prev_sibling
1135             and prevp.prev_sibling.type == token.NAME
1136             and prevp.prev_sibling.value == 'print'  # type: ignore
1137         ):
1138             # Python 2 print chevron
1139             return NO
1140
1141     elif prev.type in OPENING_BRACKETS:
1142         return NO
1143
1144     if p.type in {syms.parameters, syms.arglist}:
1145         # untyped function signatures or calls
1146         if t == token.RPAR:
1147             return NO
1148
1149         if not prev or prev.type != token.COMMA:
1150             return NO
1151
1152     elif p.type == syms.varargslist:
1153         # lambdas
1154         if t == token.RPAR:
1155             return NO
1156
1157         if prev and prev.type != token.COMMA:
1158             return NO
1159
1160     elif p.type == syms.typedargslist:
1161         # typed function signatures
1162         if not prev:
1163             return NO
1164
1165         if t == token.EQUAL:
1166             if prev.type != syms.tname:
1167                 return NO
1168
1169         elif prev.type == token.EQUAL:
1170             # A bit hacky: if the equal sign has whitespace, it means we
1171             # previously found it's a typed argument.  So, we're using that, too.
1172             return prev.prefix
1173
1174         elif prev.type != token.COMMA:
1175             return NO
1176
1177     elif p.type == syms.tname:
1178         # type names
1179         if not prev:
1180             prevp = preceding_leaf(p)
1181             if not prevp or prevp.type != token.COMMA:
1182                 return NO
1183
1184     elif p.type == syms.trailer:
1185         # attributes and calls
1186         if t == token.LPAR or t == token.RPAR:
1187             return NO
1188
1189         if not prev:
1190             if t == token.DOT:
1191                 prevp = preceding_leaf(p)
1192                 if not prevp or prevp.type != token.NUMBER:
1193                     return NO
1194
1195             elif t == token.LSQB:
1196                 return NO
1197
1198         elif prev.type != token.COMMA:
1199             return NO
1200
1201     elif p.type == syms.argument:
1202         # single argument
1203         if t == token.EQUAL:
1204             return NO
1205
1206         if not prev:
1207             prevp = preceding_leaf(p)
1208             if not prevp or prevp.type == token.LPAR:
1209                 return NO
1210
1211         elif prev.type == token.EQUAL or prev.type == token.DOUBLESTAR:
1212             return NO
1213
1214     elif p.type == syms.decorator:
1215         # decorators
1216         return NO
1217
1218     elif p.type == syms.dotted_name:
1219         if prev:
1220             return NO
1221
1222         prevp = preceding_leaf(p)
1223         if not prevp or prevp.type == token.AT or prevp.type == token.DOT:
1224             return NO
1225
1226     elif p.type == syms.classdef:
1227         if t == token.LPAR:
1228             return NO
1229
1230         if prev and prev.type == token.LPAR:
1231             return NO
1232
1233     elif p.type == syms.subscript:
1234         # indexing
1235         if not prev:
1236             assert p.parent is not None, "subscripts are always parented"
1237             if p.parent.type == syms.subscriptlist:
1238                 return SPACE
1239
1240             return NO
1241
1242         else:
1243             return NO
1244
1245     elif p.type == syms.atom:
1246         if prev and t == token.DOT:
1247             # dots, but not the first one.
1248             return NO
1249
1250     elif (
1251         p.type == syms.listmaker
1252         or p.type == syms.testlist_gexp
1253         or p.type == syms.subscriptlist
1254     ):
1255         # list interior, including unpacking
1256         if not prev:
1257             return NO
1258
1259     elif p.type == syms.dictsetmaker:
1260         # dict and set interior, including unpacking
1261         if not prev:
1262             return NO
1263
1264         if prev.type == token.DOUBLESTAR:
1265             return NO
1266
1267     elif p.type in {syms.factor, syms.star_expr}:
1268         # unary ops
1269         if not prev:
1270             prevp = preceding_leaf(p)
1271             if not prevp or prevp.type in OPENING_BRACKETS:
1272                 return NO
1273
1274             prevp_parent = prevp.parent
1275             assert prevp_parent is not None
1276             if prevp.type == token.COLON and prevp_parent.type in {
1277                 syms.subscript, syms.sliceop
1278             }:
1279                 return NO
1280
1281             elif prevp.type == token.EQUAL and prevp_parent.type == syms.argument:
1282                 return NO
1283
1284         elif t == token.NAME or t == token.NUMBER:
1285             return NO
1286
1287     elif p.type == syms.import_from:
1288         if t == token.DOT:
1289             if prev and prev.type == token.DOT:
1290                 return NO
1291
1292         elif t == token.NAME:
1293             if v == 'import':
1294                 return SPACE
1295
1296             if prev and prev.type == token.DOT:
1297                 return NO
1298
1299     elif p.type == syms.sliceop:
1300         return NO
1301
1302     return SPACE
1303
1304
1305 def preceding_leaf(node: Optional[LN]) -> Optional[Leaf]:
1306     """Returns the first leaf that precedes `node`, if any."""
1307     while node:
1308         res = node.prev_sibling
1309         if res:
1310             if isinstance(res, Leaf):
1311                 return res
1312
1313             try:
1314                 return list(res.leaves())[-1]
1315
1316             except IndexError:
1317                 return None
1318
1319         node = node.parent
1320     return None
1321
1322
1323 def is_delimiter(leaf: Leaf) -> int:
1324     """Returns the priority of the `leaf` delimiter. Returns 0 if not delimiter.
1325
1326     Higher numbers are higher priority.
1327     """
1328     if leaf.type == token.COMMA:
1329         return COMMA_PRIORITY
1330
1331     if leaf.type in COMPARATORS:
1332         return COMPARATOR_PRIORITY
1333
1334     if (
1335         leaf.type in MATH_OPERATORS
1336         and leaf.parent
1337         and leaf.parent.type not in {syms.factor, syms.star_expr}
1338     ):
1339         return MATH_PRIORITY
1340
1341     return 0
1342
1343
1344 def generate_comments(leaf: Leaf) -> Iterator[Leaf]:
1345     """Cleans the prefix of the `leaf` and generates comments from it, if any.
1346
1347     Comments in lib2to3 are shoved into the whitespace prefix.  This happens
1348     in `pgen2/driver.py:Driver.parse_tokens()`.  This was a brilliant implementation
1349     move because it does away with modifying the grammar to include all the
1350     possible places in which comments can be placed.
1351
1352     The sad consequence for us though is that comments don't "belong" anywhere.
1353     This is why this function generates simple parentless Leaf objects for
1354     comments.  We simply don't know what the correct parent should be.
1355
1356     No matter though, we can live without this.  We really only need to
1357     differentiate between inline and standalone comments.  The latter don't
1358     share the line with any code.
1359
1360     Inline comments are emitted as regular token.COMMENT leaves.  Standalone
1361     are emitted with a fake STANDALONE_COMMENT token identifier.
1362     """
1363     p = leaf.prefix
1364     if not p:
1365         return
1366
1367     if '#' not in p:
1368         return
1369
1370     consumed = 0
1371     nlines = 0
1372     for index, line in enumerate(p.split('\n')):
1373         consumed += len(line) + 1  # adding the length of the split '\n'
1374         line = line.lstrip()
1375         if not line:
1376             nlines += 1
1377         if not line.startswith('#'):
1378             continue
1379
1380         if index == 0 and leaf.type != token.ENDMARKER:
1381             comment_type = token.COMMENT  # simple trailing comment
1382         else:
1383             comment_type = STANDALONE_COMMENT
1384         comment = make_comment(line)
1385         yield Leaf(comment_type, comment, prefix='\n' * nlines)
1386
1387         if comment in {'# fmt: on', '# yapf: enable'}:
1388             raise FormatOn(consumed)
1389
1390         if comment in {'# fmt: off', '# yapf: disable'}:
1391             raise FormatOff(consumed)
1392
1393         nlines = 0
1394
1395
1396 def make_comment(content: str) -> str:
1397     """Returns a consistently formatted comment from the given `content` string.
1398
1399     All comments (except for "##", "#!", "#:") should have a single space between
1400     the hash sign and the content.
1401
1402     If `content` didn't start with a hash sign, one is provided.
1403     """
1404     content = content.rstrip()
1405     if not content:
1406         return '#'
1407
1408     if content[0] == '#':
1409         content = content[1:]
1410     if content and content[0] not in ' !:#':
1411         content = ' ' + content
1412     return '#' + content
1413
1414
1415 def split_line(
1416     line: Line, line_length: int, inner: bool = False, py36: bool = False
1417 ) -> Iterator[Line]:
1418     """Splits a `line` into potentially many lines.
1419
1420     They should fit in the allotted `line_length` but might not be able to.
1421     `inner` signifies that there were a pair of brackets somewhere around the
1422     current `line`, possibly transitively. This means we can fallback to splitting
1423     by delimiters if the LHS/RHS don't yield any results.
1424
1425     If `py36` is True, splitting may generate syntax that is only compatible
1426     with Python 3.6 and later.
1427     """
1428     if isinstance(line, UnformattedLines):
1429         yield line
1430         return
1431
1432     line_str = str(line).strip('\n')
1433     if len(line_str) <= line_length and '\n' not in line_str:
1434         yield line
1435         return
1436
1437     if line.is_def:
1438         split_funcs = [left_hand_split]
1439     elif line.inside_brackets:
1440         split_funcs = [delimiter_split]
1441         if '\n' not in line_str:
1442             # Only attempt RHS if we don't have multiline strings or comments
1443             # on this line.
1444             split_funcs.append(right_hand_split)
1445     else:
1446         split_funcs = [right_hand_split]
1447     for split_func in split_funcs:
1448         # We are accumulating lines in `result` because we might want to abort
1449         # mission and return the original line in the end, or attempt a different
1450         # split altogether.
1451         result: List[Line] = []
1452         try:
1453             for l in split_func(line, py36=py36):
1454                 if str(l).strip('\n') == line_str:
1455                     raise CannotSplit("Split function returned an unchanged result")
1456
1457                 result.extend(
1458                     split_line(l, line_length=line_length, inner=True, py36=py36)
1459                 )
1460         except CannotSplit as cs:
1461             continue
1462
1463         else:
1464             yield from result
1465             break
1466
1467     else:
1468         yield line
1469
1470
1471 def left_hand_split(line: Line, py36: bool = False) -> Iterator[Line]:
1472     """Splits line into many lines, starting with the first matching bracket pair.
1473
1474     Note: this usually looks weird, only use this for function definitions.
1475     Prefer RHS otherwise.
1476     """
1477     head = Line(depth=line.depth)
1478     body = Line(depth=line.depth + 1, inside_brackets=True)
1479     tail = Line(depth=line.depth)
1480     tail_leaves: List[Leaf] = []
1481     body_leaves: List[Leaf] = []
1482     head_leaves: List[Leaf] = []
1483     current_leaves = head_leaves
1484     matching_bracket = None
1485     for leaf in line.leaves:
1486         if (
1487             current_leaves is body_leaves
1488             and leaf.type in CLOSING_BRACKETS
1489             and leaf.opening_bracket is matching_bracket
1490         ):
1491             current_leaves = tail_leaves if body_leaves else head_leaves
1492         current_leaves.append(leaf)
1493         if current_leaves is head_leaves:
1494             if leaf.type in OPENING_BRACKETS:
1495                 matching_bracket = leaf
1496                 current_leaves = body_leaves
1497     # Since body is a new indent level, remove spurious leading whitespace.
1498     if body_leaves:
1499         normalize_prefix(body_leaves[0], inside_brackets=True)
1500     # Build the new lines.
1501     for result, leaves in (
1502         (head, head_leaves), (body, body_leaves), (tail, tail_leaves)
1503     ):
1504         for leaf in leaves:
1505             result.append(leaf, preformatted=True)
1506             comment_after = line.comments.get(id(leaf))
1507             if comment_after:
1508                 result.append(comment_after, preformatted=True)
1509     bracket_split_succeeded_or_raise(head, body, tail)
1510     for result in (head, body, tail):
1511         if result:
1512             yield result
1513
1514
1515 def right_hand_split(line: Line, py36: bool = False) -> Iterator[Line]:
1516     """Splits line into many lines, starting with the last matching bracket pair."""
1517     head = Line(depth=line.depth)
1518     body = Line(depth=line.depth + 1, inside_brackets=True)
1519     tail = Line(depth=line.depth)
1520     tail_leaves: List[Leaf] = []
1521     body_leaves: List[Leaf] = []
1522     head_leaves: List[Leaf] = []
1523     current_leaves = tail_leaves
1524     opening_bracket = None
1525     for leaf in reversed(line.leaves):
1526         if current_leaves is body_leaves:
1527             if leaf is opening_bracket:
1528                 current_leaves = head_leaves if body_leaves else tail_leaves
1529         current_leaves.append(leaf)
1530         if current_leaves is tail_leaves:
1531             if leaf.type in CLOSING_BRACKETS:
1532                 opening_bracket = leaf.opening_bracket
1533                 current_leaves = body_leaves
1534     tail_leaves.reverse()
1535     body_leaves.reverse()
1536     head_leaves.reverse()
1537     # Since body is a new indent level, remove spurious leading whitespace.
1538     if body_leaves:
1539         normalize_prefix(body_leaves[0], inside_brackets=True)
1540     # Build the new lines.
1541     for result, leaves in (
1542         (head, head_leaves), (body, body_leaves), (tail, tail_leaves)
1543     ):
1544         for leaf in leaves:
1545             result.append(leaf, preformatted=True)
1546             comment_after = line.comments.get(id(leaf))
1547             if comment_after:
1548                 result.append(comment_after, preformatted=True)
1549     bracket_split_succeeded_or_raise(head, body, tail)
1550     for result in (head, body, tail):
1551         if result:
1552             yield result
1553
1554
1555 def bracket_split_succeeded_or_raise(head: Line, body: Line, tail: Line) -> None:
1556     """Raise :exc:`CannotSplit` if the last left- or right-hand split failed.
1557
1558     Do nothing otherwise.
1559
1560     A left- or right-hand split is based on a pair of brackets. Content before
1561     (and including) the opening bracket is left on one line, content inside the
1562     brackets is put on a separate line, and finally content starting with and
1563     following the closing bracket is put on a separate line.
1564
1565     Those are called `head`, `body`, and `tail`, respectively. If the split
1566     produced the same line (all content in `head`) or ended up with an empty `body`
1567     and the `tail` is just the closing bracket, then it's considered failed.
1568     """
1569     tail_len = len(str(tail).strip())
1570     if not body:
1571         if tail_len == 0:
1572             raise CannotSplit("Splitting brackets produced the same line")
1573
1574         elif tail_len < 3:
1575             raise CannotSplit(
1576                 f"Splitting brackets on an empty body to save "
1577                 f"{tail_len} characters is not worth it"
1578             )
1579
1580
1581 def delimiter_split(line: Line, py36: bool = False) -> Iterator[Line]:
1582     """Splits according to delimiters of the highest priority.
1583
1584     This kind of split doesn't increase indentation.
1585     If `py36` is True, the split will add trailing commas also in function
1586     signatures that contain `*` and `**`.
1587     """
1588     try:
1589         last_leaf = line.leaves[-1]
1590     except IndexError:
1591         raise CannotSplit("Line empty")
1592
1593     delimiters = line.bracket_tracker.delimiters
1594     try:
1595         delimiter_priority = line.bracket_tracker.max_delimiter_priority(
1596             exclude={id(last_leaf)}
1597         )
1598     except ValueError:
1599         raise CannotSplit("No delimiters found")
1600
1601     current_line = Line(depth=line.depth, inside_brackets=line.inside_brackets)
1602     lowest_depth = sys.maxsize
1603     trailing_comma_safe = True
1604     for leaf in line.leaves:
1605         current_line.append(leaf, preformatted=True)
1606         comment_after = line.comments.get(id(leaf))
1607         if comment_after:
1608             current_line.append(comment_after, preformatted=True)
1609         lowest_depth = min(lowest_depth, leaf.bracket_depth)
1610         if (
1611             leaf.bracket_depth == lowest_depth
1612             and leaf.type == token.STAR
1613             or leaf.type == token.DOUBLESTAR
1614         ):
1615             trailing_comma_safe = trailing_comma_safe and py36
1616         leaf_priority = delimiters.get(id(leaf))
1617         if leaf_priority == delimiter_priority:
1618             normalize_prefix(current_line.leaves[0], inside_brackets=True)
1619             yield current_line
1620
1621             current_line = Line(depth=line.depth, inside_brackets=line.inside_brackets)
1622     if current_line:
1623         if (
1624             delimiter_priority == COMMA_PRIORITY
1625             and current_line.leaves[-1].type != token.COMMA
1626             and trailing_comma_safe
1627         ):
1628             current_line.append(Leaf(token.COMMA, ','))
1629         normalize_prefix(current_line.leaves[0], inside_brackets=True)
1630         yield current_line
1631
1632
1633 def is_import(leaf: Leaf) -> bool:
1634     """Returns True if the given leaf starts an import statement."""
1635     p = leaf.parent
1636     t = leaf.type
1637     v = leaf.value
1638     return bool(
1639         t == token.NAME
1640         and (
1641             (v == 'import' and p and p.type == syms.import_name)
1642             or (v == 'from' and p and p.type == syms.import_from)
1643         )
1644     )
1645
1646
1647 def normalize_prefix(leaf: Leaf, *, inside_brackets: bool) -> None:
1648     """Leaves existing extra newlines if not `inside_brackets`.
1649
1650     Removes everything else.  Note: don't use backslashes for formatting or
1651     you'll lose your voting rights.
1652     """
1653     if not inside_brackets:
1654         spl = leaf.prefix.split('#')
1655         if '\\' not in spl[0]:
1656             nl_count = spl[-1].count('\n')
1657             if len(spl) > 1:
1658                 nl_count -= 1
1659             leaf.prefix = '\n' * nl_count
1660             return
1661
1662     leaf.prefix = ''
1663
1664
1665 def is_python36(node: Node) -> bool:
1666     """Returns True if the current file is using Python 3.6+ features.
1667
1668     Currently looking for:
1669     - f-strings; and
1670     - trailing commas after * or ** in function signatures.
1671     """
1672     for n in node.pre_order():
1673         if n.type == token.STRING:
1674             value_head = n.value[:2]  # type: ignore
1675             if value_head in {'f"', 'F"', "f'", "F'", 'rf', 'fr', 'RF', 'FR'}:
1676                 return True
1677
1678         elif (
1679             n.type == syms.typedargslist
1680             and n.children
1681             and n.children[-1].type == token.COMMA
1682         ):
1683             for ch in n.children:
1684                 if ch.type == token.STAR or ch.type == token.DOUBLESTAR:
1685                     return True
1686
1687     return False
1688
1689
1690 PYTHON_EXTENSIONS = {'.py'}
1691 BLACKLISTED_DIRECTORIES = {
1692     'build', 'buck-out', 'dist', '_build', '.git', '.hg', '.mypy_cache', '.tox', '.venv'
1693 }
1694
1695
1696 def gen_python_files_in_dir(path: Path) -> Iterator[Path]:
1697     """Generates all files under `path` which aren't under BLACKLISTED_DIRECTORIES
1698     and have one of the PYTHON_EXTENSIONS.
1699     """
1700     for child in path.iterdir():
1701         if child.is_dir():
1702             if child.name in BLACKLISTED_DIRECTORIES:
1703                 continue
1704
1705             yield from gen_python_files_in_dir(child)
1706
1707         elif child.suffix in PYTHON_EXTENSIONS:
1708             yield child
1709
1710
1711 @dataclass
1712 class Report:
1713     """Provides a reformatting counter. Can be rendered with `str(report)`."""
1714     check: bool = False
1715     change_count: int = 0
1716     same_count: int = 0
1717     failure_count: int = 0
1718
1719     def done(self, src: Path, changed: bool) -> None:
1720         """Increment the counter for successful reformatting. Write out a message."""
1721         if changed:
1722             reformatted = 'would reformat' if self.check else 'reformatted'
1723             out(f'{reformatted} {src}')
1724             self.change_count += 1
1725         else:
1726             out(f'{src} already well formatted, good job.', bold=False)
1727             self.same_count += 1
1728
1729     def failed(self, src: Path, message: str) -> None:
1730         """Increment the counter for failed reformatting. Write out a message."""
1731         err(f'error: cannot format {src}: {message}')
1732         self.failure_count += 1
1733
1734     @property
1735     def return_code(self) -> int:
1736         """Which return code should the app use considering the current state."""
1737         # According to http://tldp.org/LDP/abs/html/exitcodes.html starting with
1738         # 126 we have special returncodes reserved by the shell.
1739         if self.failure_count:
1740             return 123
1741
1742         elif self.change_count and self.check:
1743             return 1
1744
1745         return 0
1746
1747     def __str__(self) -> str:
1748         """A color report of the current state.
1749
1750         Use `click.unstyle` to remove colors.
1751         """
1752         if self.check:
1753             reformatted = "would be reformatted"
1754             unchanged = "would be left unchanged"
1755             failed = "would fail to reformat"
1756         else:
1757             reformatted = "reformatted"
1758             unchanged = "left unchanged"
1759             failed = "failed to reformat"
1760         report = []
1761         if self.change_count:
1762             s = 's' if self.change_count > 1 else ''
1763             report.append(
1764                 click.style(f'{self.change_count} file{s} {reformatted}', bold=True)
1765             )
1766         if self.same_count:
1767             s = 's' if self.same_count > 1 else ''
1768             report.append(f'{self.same_count} file{s} {unchanged}')
1769         if self.failure_count:
1770             s = 's' if self.failure_count > 1 else ''
1771             report.append(
1772                 click.style(f'{self.failure_count} file{s} {failed}', fg='red')
1773             )
1774         return ', '.join(report) + '.'
1775
1776
1777 def assert_equivalent(src: str, dst: str) -> None:
1778     """Raises AssertionError if `src` and `dst` aren't equivalent.
1779
1780     This is a temporary sanity check until Black becomes stable.
1781     """
1782
1783     import ast
1784     import traceback
1785
1786     def _v(node: ast.AST, depth: int = 0) -> Iterator[str]:
1787         """Simple visitor generating strings to compare ASTs by content."""
1788         yield f"{'  ' * depth}{node.__class__.__name__}("
1789
1790         for field in sorted(node._fields):
1791             try:
1792                 value = getattr(node, field)
1793             except AttributeError:
1794                 continue
1795
1796             yield f"{'  ' * (depth+1)}{field}="
1797
1798             if isinstance(value, list):
1799                 for item in value:
1800                     if isinstance(item, ast.AST):
1801                         yield from _v(item, depth + 2)
1802
1803             elif isinstance(value, ast.AST):
1804                 yield from _v(value, depth + 2)
1805
1806             else:
1807                 yield f"{'  ' * (depth+2)}{value!r},  # {value.__class__.__name__}"
1808
1809         yield f"{'  ' * depth})  # /{node.__class__.__name__}"
1810
1811     try:
1812         src_ast = ast.parse(src)
1813     except Exception as exc:
1814         major, minor = sys.version_info[:2]
1815         raise AssertionError(
1816             f"cannot use --safe with this file; failed to parse source file "
1817             f"with Python {major}.{minor}'s builtin AST. Re-run with --fast "
1818             f"or stop using deprecated Python 2 syntax. AST error message: {exc}"
1819         )
1820
1821     try:
1822         dst_ast = ast.parse(dst)
1823     except Exception as exc:
1824         log = dump_to_file(''.join(traceback.format_tb(exc.__traceback__)), dst)
1825         raise AssertionError(
1826             f"INTERNAL ERROR: Black produced invalid code: {exc}. "
1827             f"Please report a bug on https://github.com/ambv/black/issues.  "
1828             f"This invalid output might be helpful: {log}"
1829         ) from None
1830
1831     src_ast_str = '\n'.join(_v(src_ast))
1832     dst_ast_str = '\n'.join(_v(dst_ast))
1833     if src_ast_str != dst_ast_str:
1834         log = dump_to_file(diff(src_ast_str, dst_ast_str, 'src', 'dst'))
1835         raise AssertionError(
1836             f"INTERNAL ERROR: Black produced code that is not equivalent to "
1837             f"the source.  "
1838             f"Please report a bug on https://github.com/ambv/black/issues.  "
1839             f"This diff might be helpful: {log}"
1840         ) from None
1841
1842
1843 def assert_stable(src: str, dst: str, line_length: int) -> None:
1844     """Raises AssertionError if `dst` reformats differently the second time.
1845
1846     This is a temporary sanity check until Black becomes stable.
1847     """
1848     newdst = format_str(dst, line_length=line_length)
1849     if dst != newdst:
1850         log = dump_to_file(
1851             diff(src, dst, 'source', 'first pass'),
1852             diff(dst, newdst, 'first pass', 'second pass'),
1853         )
1854         raise AssertionError(
1855             f"INTERNAL ERROR: Black produced different code on the second pass "
1856             f"of the formatter.  "
1857             f"Please report a bug on https://github.com/ambv/black/issues.  "
1858             f"This diff might be helpful: {log}"
1859         ) from None
1860
1861
1862 def dump_to_file(*output: str) -> str:
1863     """Dumps `output` to a temporary file. Returns path to the file."""
1864     import tempfile
1865
1866     with tempfile.NamedTemporaryFile(
1867         mode='w', prefix='blk_', suffix='.log', delete=False
1868     ) as f:
1869         for lines in output:
1870             f.write(lines)
1871             f.write('\n')
1872     return f.name
1873
1874
1875 def diff(a: str, b: str, a_name: str, b_name: str) -> str:
1876     """Returns a udiff string between strings `a` and `b`."""
1877     import difflib
1878
1879     a_lines = [line + '\n' for line in a.split('\n')]
1880     b_lines = [line + '\n' for line in b.split('\n')]
1881     return ''.join(
1882         difflib.unified_diff(a_lines, b_lines, fromfile=a_name, tofile=b_name, n=5)
1883     )
1884
1885
1886 if __name__ == '__main__':
1887     main()