[etc/vim.git] / black.py

#!/usr/bin/env python3
import asyncio
from asyncio.base_events import BaseEventLoop
from concurrent.futures import Executor, ProcessPoolExecutor
from functools import partial
import keyword
import os
from pathlib import Path
import tokenize
from typing import (
    Dict, Generic, Iterable, Iterator, List, Optional, Set, Tuple, TypeVar, Union
)

from attr import attrib, dataclass, Factory
import click

# lib2to3 fork
from blib2to3.pytree import Node, Leaf, type_repr
from blib2to3 import pygram, pytree
from blib2to3.pgen2 import driver, token
from blib2to3.pgen2.parse import ParseError

__version__ = "18.3a1"
DEFAULT_LINE_LENGTH = 88
# types
syms = pygram.python_symbols
FileContent = str
Encoding = str
Depth = int
NodeType = int
LeafID = int
Priority = int
LN = Union[Leaf, Node]
out = partial(click.secho, bold=True, err=True)
err = partial(click.secho, fg='red', err=True)


class NothingChanged(UserWarning):
    """Raised by `format_file` when the reformatted code is the same as source."""


class CannotSplit(Exception):
    """A readable split that fits the allotted line length is impossible.

    Raised by `left_hand_split()` and `right_hand_split()`.
    """


@click.command()
@click.option(
    '-l',
    '--line-length',
    type=int,
    default=DEFAULT_LINE_LENGTH,
    help='How many character per line to allow.',
    show_default=True,
)
@click.option(
    '--fast/--safe',
    is_flag=True,
    help='If --fast given, skip temporary sanity checks. [default: --safe]',
)
@click.version_option(version=__version__)
@click.argument(
    'src',
    nargs=-1,
    type=click.Path(exists=True, file_okay=True, dir_okay=True, readable=True),
)
@click.pass_context
def main(ctx: click.Context, line_length: int, fast: bool, src: List[str]) -> None:
    """The uncompromising code formatter."""
    sources: List[Path] = []
    for s in src:
        p = Path(s)
        if p.is_dir():
            sources.extend(gen_python_files_in_dir(p))
        elif p.is_file():
            # if a file was explicitly given, we don't care about its extension
            sources.append(p)
        else:
            err(f'invalid path: {s}')
    if len(sources) == 0:
        ctx.exit(0)
    elif len(sources) == 1:
        p = sources[0]
        report = Report()
        try:
            changed = format_file_in_place(p, line_length=line_length, fast=fast)
            report.done(p, changed)
        except Exception as exc:
            report.failed(p, str(exc))
        ctx.exit(report.return_code)
    else:
        loop = asyncio.get_event_loop()
        executor = ProcessPoolExecutor(max_workers=os.cpu_count())
        return_code = 1
        try:
            return_code = loop.run_until_complete(
                schedule_formatting(sources, line_length, fast, loop, executor)
            )
        finally:
            loop.close()
            ctx.exit(return_code)


async def schedule_formatting(
    sources: List[Path],
    line_length: int,
    fast: bool,
    loop: BaseEventLoop,
    executor: Executor,
) -> int:
    tasks = {
        src: loop.run_in_executor(
            executor, format_file_in_place, src, line_length, fast
        )
        for src in sources
    }
    await asyncio.wait(tasks.values())
    cancelled = []
    report = Report()
    for src, task in tasks.items():
        if not task.done():
            report.failed(src, 'timed out, cancelling')
            task.cancel()
            cancelled.append(task)
        elif task.exception():
            report.failed(src, str(task.exception()))
        else:
            report.done(src, task.result())
    if cancelled:
        await asyncio.wait(cancelled, timeout=2)
    out('All done! ✨ 🍰 ✨')
    click.echo(str(report))
    return report.return_code


def format_file_in_place(src: Path, line_length: int, fast: bool) -> bool:
    """Format the file and rewrite if changed. Return True if changed."""
    try:
        contents, encoding = format_file(src, line_length=line_length, fast=fast)
    except NothingChanged:
        return False

    with open(src, "w", encoding=encoding) as f:
        f.write(contents)
    return True


def format_file(
    src: Path, line_length: int, fast: bool
) -> Tuple[FileContent, Encoding]:
    """Reformats a file and returns its contents and encoding."""
    with tokenize.open(src) as src_buffer:
        src_contents = src_buffer.read()
    if src_contents.strip() == '':
        raise NothingChanged(src)

    dst_contents = format_str(src_contents, line_length=line_length)
    if src_contents == dst_contents:
        raise NothingChanged(src)

    if not fast:
        assert_equivalent(src_contents, dst_contents)
        assert_stable(src_contents, dst_contents, line_length=line_length)
    return dst_contents, src_buffer.encoding


def format_str(src_contents: str, line_length: int) -> FileContent:
    """Reformats a string and returns new contents."""
    src_node = lib2to3_parse(src_contents)
    dst_contents = ""
    comments: List[Line] = []
    lines = LineGenerator()
    elt = EmptyLineTracker()
    empty_line = Line()
    after = 0
    for current_line in lines.visit(src_node):
        for _ in range(after):
            dst_contents += str(empty_line)
        before, after = elt.maybe_empty_lines(current_line)
        for _ in range(before):
            dst_contents += str(empty_line)
        if not current_line.is_comment:
            for comment in comments:
                dst_contents += str(comment)
            comments = []
            for line in split_line(current_line, line_length=line_length):
                dst_contents += str(line)
        else:
            comments.append(current_line)
    for comment in comments:
        dst_contents += str(comment)
    return dst_contents


def lib2to3_parse(src_txt: str) -> Node:
    """Given a string with source, return the lib2to3 Node."""
    grammar = pygram.python_grammar_no_print_statement
    drv = driver.Driver(grammar, pytree.convert)
    if src_txt[-1] != '\n':
        nl = '\r\n' if '\r\n' in src_txt[:1024] else '\n'
        src_txt += nl
    try:
        result = drv.parse_string(src_txt, True)
    except ParseError as pe:
        lineno, column = pe.context[1]
        lines = src_txt.splitlines()
        try:
            faulty_line = lines[lineno - 1]
        except IndexError:
            faulty_line = "<line number missing in source>"
        raise ValueError(f"Cannot parse: {lineno}:{column}: {faulty_line}") from None

    if isinstance(result, Leaf):
        result = Node(syms.file_input, [result])
    return result


def lib2to3_unparse(node: Node) -> str:
    """Given a lib2to3 node, return its string representation."""
    code = str(node)
    return code


T = TypeVar('T')


class Visitor(Generic[T]):
    """Basic lib2to3 visitor that yields things on visiting."""

    def visit(self, node: LN) -> Iterator[T]:
        if node.type < 256:
            name = token.tok_name[node.type]
        else:
            name = type_repr(node.type)
        yield from getattr(self, f'visit_{name}', self.visit_default)(node)

    def visit_default(self, node: LN) -> Iterator[T]:
        if isinstance(node, Node):
            for child in node.children:
                yield from self.visit(child)


@dataclass
class DebugVisitor(Visitor[T]):
    tree_depth: int = attrib(default=0)

    def visit_default(self, node: LN) -> Iterator[T]:
        indent = ' ' * (2 * self.tree_depth)
        if isinstance(node, Node):
            _type = type_repr(node.type)
            out(f'{indent}{_type}', fg='yellow')
            self.tree_depth += 1
            for child in node.children:
                yield from self.visit(child)

            self.tree_depth -= 1
            out(f'{indent}/{_type}', fg='yellow', bold=False)
        else:
            _type = token.tok_name.get(node.type, str(node.type))
            out(f'{indent}{_type}', fg='blue', nl=False)
            if node.prefix:
                # We don't have to handle prefixes for `Node` objects since
                # that delegates to the first child anyway.
                out(f' {node.prefix!r}', fg='green', bold=False, nl=False)
            out(f' {node.value!r}', fg='blue', bold=False)


KEYWORDS = set(keyword.kwlist)
WHITESPACE = {token.DEDENT, token.INDENT, token.NEWLINE}
FLOW_CONTROL = {'return', 'raise', 'break', 'continue'}
STATEMENT = {
    syms.if_stmt,
    syms.while_stmt,
    syms.for_stmt,
    syms.try_stmt,
    syms.except_clause,
    syms.with_stmt,
    syms.funcdef,
    syms.classdef,
}
STANDALONE_COMMENT = 153
LOGIC_OPERATORS = {'and', 'or'}
COMPARATORS = {
    token.LESS,
    token.GREATER,
    token.EQEQUAL,
    token.NOTEQUAL,
    token.LESSEQUAL,
    token.GREATEREQUAL,
}
MATH_OPERATORS = {
    token.PLUS,
    token.MINUS,
    token.STAR,
    token.SLASH,
    token.VBAR,
    token.AMPER,
    token.PERCENT,
    token.CIRCUMFLEX,
    token.LEFTSHIFT,
    token.RIGHTSHIFT,
    token.DOUBLESTAR,
    token.DOUBLESLASH,
}
COMPREHENSION_PRIORITY = 20
COMMA_PRIORITY = 10
LOGIC_PRIORITY = 5
STRING_PRIORITY = 4
COMPARATOR_PRIORITY = 3
MATH_PRIORITY = 1


@dataclass
class BracketTracker:
    depth: int = attrib(default=0)
    bracket_match: Dict[Tuple[Depth, NodeType], Leaf] = attrib(default=Factory(dict))
    delimiters: Dict[LeafID, Priority] = attrib(default=Factory(dict))
    previous: Optional[Leaf] = attrib(default=None)

    def mark(self, leaf: Leaf) -> None:
        if leaf.type == token.COMMENT:
            return

        if leaf.type in CLOSING_BRACKETS:
            self.depth -= 1
            opening_bracket = self.bracket_match.pop((self.depth, leaf.type))
            leaf.opening_bracket = opening_bracket  # type: ignore
        leaf.bracket_depth = self.depth  # type: ignore
        if self.depth == 0:
            delim = is_delimiter(leaf)
            if delim:
                self.delimiters[id(leaf)] = delim
            elif self.previous is not None:
                if leaf.type == token.STRING and self.previous.type == token.STRING:
                    self.delimiters[id(self.previous)] = STRING_PRIORITY
                elif (
                    leaf.type == token.NAME and
                    leaf.value == 'for' and
                    leaf.parent and
                    leaf.parent.type in {syms.comp_for, syms.old_comp_for}
                ):
                    self.delimiters[id(self.previous)] = COMPREHENSION_PRIORITY
                elif (
                    leaf.type == token.NAME and
                    leaf.value == 'if' and
                    leaf.parent and
                    leaf.parent.type in {syms.comp_if, syms.old_comp_if}
                ):
                    self.delimiters[id(self.previous)] = COMPREHENSION_PRIORITY
        if leaf.type in OPENING_BRACKETS:
            self.bracket_match[self.depth, BRACKET[leaf.type]] = leaf
            self.depth += 1
        self.previous = leaf

    def any_open_brackets(self) -> bool:
        """Returns True if there is an yet unmatched open bracket on the line."""
        return bool(self.bracket_match)

    def max_priority(self, exclude: Iterable[LeafID] = ()) -> int:
        """Returns the highest priority of a delimiter found on the line.

        Values are consistent with what `is_delimiter()` returns.
        """
        return max(v for k, v in self.delimiters.items() if k not in exclude)


@dataclass
class Line:
    depth: int = attrib(default=0)
    leaves: List[Leaf] = attrib(default=Factory(list))
    comments: Dict[LeafID, Leaf] = attrib(default=Factory(dict))
    bracket_tracker: BracketTracker = attrib(default=Factory(BracketTracker))
    inside_brackets: bool = attrib(default=False)

    def append(self, leaf: Leaf, preformatted: bool = False) -> None:
        has_value = leaf.value.strip()
        if not has_value:
            return

        if self.leaves and not preformatted:
            # Note: at this point leaf.prefix should be empty except for
            # imports, for which we only preserve newlines.
            leaf.prefix += whitespace(leaf)
        if self.inside_brackets or not preformatted:
            self.bracket_tracker.mark(leaf)
            self.maybe_remove_trailing_comma(leaf)
            if self.maybe_adapt_standalone_comment(leaf):
                return

        if not self.append_comment(leaf):
            self.leaves.append(leaf)

    @property
    def is_comment(self) -> bool:
        return bool(self) and self.leaves[0].type == STANDALONE_COMMENT

    @property
    def is_decorator(self) -> bool:
        return bool(self) and self.leaves[0].type == token.AT

    @property
    def is_import(self) -> bool:
        return bool(self) and is_import(self.leaves[0])

    @property
    def is_class(self) -> bool:
        return (
            bool(self) and
            self.leaves[0].type == token.NAME and
            self.leaves[0].value == 'class'
        )

    @property
    def is_def(self) -> bool:
        """Also returns True for async defs."""
        try:
            first_leaf = self.leaves[0]
        except IndexError:
            return False

        try:
            second_leaf: Optional[Leaf] = self.leaves[1]
        except IndexError:
            second_leaf = None
        return (
            (first_leaf.type == token.NAME and first_leaf.value == 'def') or
            (
                first_leaf.type == token.NAME and
                first_leaf.value == 'async' and
                second_leaf is not None and
                second_leaf.type == token.NAME and
                second_leaf.value == 'def'
            )
        )

    @property
    def is_flow_control(self) -> bool:
        return (
            bool(self) and
            self.leaves[0].type == token.NAME and
            self.leaves[0].value in FLOW_CONTROL
        )

    @property
    def is_yield(self) -> bool:
        return (
            bool(self) and
            self.leaves[0].type == token.NAME and
            self.leaves[0].value == 'yield'
        )

    def maybe_remove_trailing_comma(self, closing: Leaf) -> bool:
        if not (
            self.leaves and
            self.leaves[-1].type == token.COMMA and
            closing.type in CLOSING_BRACKETS
        ):
            return False

        if closing.type == token.RSQB or closing.type == token.RBRACE:
            self.leaves.pop()
            return True

        # For parens let's check if it's safe to remove the comma.  If the
        # trailing one is the only one, we might mistakenly change a tuple
        # into a different type by removing the comma.
        depth = closing.bracket_depth + 1  # type: ignore
        commas = 0
        opening = closing.opening_bracket  # type: ignore
        for _opening_index, leaf in enumerate(self.leaves):
            if leaf is opening:
                break

        else:
            return False

        for leaf in self.leaves[_opening_index + 1:]:
            if leaf is closing:
                break

            bracket_depth = leaf.bracket_depth  # type: ignore
            if bracket_depth == depth and leaf.type == token.COMMA:
                commas += 1
        if commas > 1:
            self.leaves.pop()
            return True

        return False

    def maybe_adapt_standalone_comment(self, comment: Leaf) -> bool:
        """Hack a standalone comment to act as a trailing comment for line splitting.

        If this line has brackets and a standalone `comment`, we need to adapt
        it to be able to still reformat the line.

        This is not perfect, the line to which the standalone comment gets
        appended will appear "too long" when splitting.
        """
        if not (
            comment.type == STANDALONE_COMMENT and
            self.bracket_tracker.any_open_brackets()
        ):
            return False

        comment.type = token.COMMENT
        comment.prefix = '\n' + '    ' * (self.depth + 1)
        return self.append_comment(comment)

    def append_comment(self, comment: Leaf) -> bool:
        if comment.type != token.COMMENT:
            return False

        try:
            after = id(self.last_non_delimiter())
        except LookupError:
            comment.type = STANDALONE_COMMENT
            comment.prefix = ''
            return False

        else:
            if after in self.comments:
                self.comments[after].value += str(comment)
            else:
                self.comments[after] = comment
            return True

    def last_non_delimiter(self) -> Leaf:
        for i in range(len(self.leaves)):
            last = self.leaves[-i - 1]
            if not is_delimiter(last):
                return last

        raise LookupError("No non-delimiters found")

    def __str__(self) -> str:
        if not self:
            return '\n'

        indent = '    ' * self.depth
        leaves = iter(self.leaves)
        first = next(leaves)
        res = f'{first.prefix}{indent}{first.value}'
        for leaf in leaves:
            res += str(leaf)
        for comment in self.comments.values():
            res += str(comment)
        return res + '\n'

    def __bool__(self) -> bool:
        return bool(self.leaves or self.comments)


@dataclass
class EmptyLineTracker:
    """Provides a stateful method that returns the number of potential extra
    empty lines needed before and after the currently processed line.

    Note: this tracker works on lines that haven't been split yet.
    """
    previous_line: Optional[Line] = attrib(default=None)
    previous_after: int = attrib(default=0)
    previous_defs: List[int] = attrib(default=Factory(list))

    def maybe_empty_lines(self, current_line: Line) -> Tuple[int, int]:
        """Returns the number of extra empty lines before and after the `current_line`.

        This is for separating `def`, `async def` and `class` with extra empty lines
        (two on module-level), as well as providing an extra empty line after flow
        control keywords to make them more prominent.
        """
        before, after = self._maybe_empty_lines(current_line)
        self.previous_after = after
        self.previous_line = current_line
        return before, after

    def _maybe_empty_lines(self, current_line: Line) -> Tuple[int, int]:
        before = 0
        depth = current_line.depth
        while self.previous_defs and self.previous_defs[-1] >= depth:
            self.previous_defs.pop()
            before = (1 if depth else 2) - self.previous_after
        is_decorator = current_line.is_decorator
        if is_decorator or current_line.is_def or current_line.is_class:
            if not is_decorator:
                self.previous_defs.append(depth)
            if self.previous_line is None:
                # Don't insert empty lines before the first line in the file.
                return 0, 0

            if self.previous_line and self.previous_line.is_decorator:
                # Don't insert empty lines between decorators.
                return 0, 0

            newlines = 2
            if current_line.depth:
                newlines -= 1
            newlines -= self.previous_after
            return newlines, 0

        if current_line.is_flow_control:
            return before, 1

        if (
            self.previous_line and
            self.previous_line.is_import and
            not current_line.is_import and
            depth == self.previous_line.depth
        ):
            return (before or 1), 0

        if (
            self.previous_line and
            self.previous_line.is_yield and
            (not current_line.is_yield or depth != self.previous_line.depth)
        ):
            return (before or 1), 0

        return before, 0


@dataclass
class LineGenerator(Visitor[Line]):
    """Generates reformatted Line objects.  Empty lines are not emitted.

    Note: destroys the tree it's visiting by mutating prefixes of its leaves
    in ways that will no longer stringify to valid Python code on the tree.
    """
    current_line: Line = attrib(default=Factory(Line))
    standalone_comments: List[Leaf] = attrib(default=Factory(list))

    def line(self, indent: int = 0) -> Iterator[Line]:
        """Generate a line.

        If the line is empty, only emit if it makes sense.
        If the line is too long, split it first and then generate.

        If any lines were generated, set up a new current_line.
        """
        if not self.current_line:
            self.current_line.depth += indent
            return  # Line is empty, don't emit. Creating a new one unnecessary.

        complete_line = self.current_line
        self.current_line = Line(depth=complete_line.depth + indent)
        yield complete_line

    def visit_default(self, node: LN) -> Iterator[Line]:
        if isinstance(node, Leaf):
            for comment in generate_comments(node):
                if self.current_line.bracket_tracker.any_open_brackets():
                    # any comment within brackets is subject to splitting
                    self.current_line.append(comment)
                elif comment.type == token.COMMENT:
                    # regular trailing comment
                    self.current_line.append(comment)
                    yield from self.line()

                else:
                    # regular standalone comment, to be processed later (see
                    # docstring in `generate_comments()`
                    self.standalone_comments.append(comment)
            normalize_prefix(node)
            if node.type not in WHITESPACE:
                for comment in self.standalone_comments:
                    yield from self.line()

                    self.current_line.append(comment)
                    yield from self.line()

                self.standalone_comments = []
                self.current_line.append(node)
        yield from super().visit_default(node)

    def visit_suite(self, node: Node) -> Iterator[Line]:
        """Body of a statement after a colon."""
        children = iter(node.children)
        # Process newline before indenting.  It might contain an inline
        # comment that should go right after the colon.
        newline = next(children)
        yield from self.visit(newline)
        yield from self.line(+1)

        for child in children:
            yield from self.visit(child)

        yield from self.line(-1)

    def visit_stmt(self, node: Node, keywords: Set[str]) -> Iterator[Line]:
        """Visit a statement.

        The relevant Python language keywords for this statement are NAME leaves
        within it.
        """
        for child in node.children:
            if child.type == token.NAME and child.value in keywords:  # type: ignore
                yield from self.line()

            yield from self.visit(child)

    def visit_simple_stmt(self, node: Node) -> Iterator[Line]:
        """A statement without nested statements."""
        is_suite_like = node.parent and node.parent.type in STATEMENT
        if is_suite_like:
            yield from self.line(+1)
            yield from self.visit_default(node)
            yield from self.line(-1)

        else:
            yield from self.line()
            yield from self.visit_default(node)

    def visit_async_stmt(self, node: Node) -> Iterator[Line]:
        yield from self.line()

        children = iter(node.children)
        for child in children:
            yield from self.visit(child)

            if child.type == token.NAME and child.value == 'async':  # type: ignore
                break

        internal_stmt = next(children)
        for child in internal_stmt.children:
            yield from self.visit(child)

    def visit_decorators(self, node: Node) -> Iterator[Line]:
        for child in node.children:
            yield from self.line()
            yield from self.visit(child)

    def visit_SEMI(self, leaf: Leaf) -> Iterator[Line]:
        yield from self.line()

    def visit_ENDMARKER(self, leaf: Leaf) -> Iterator[Line]:
        yield from self.visit_default(leaf)
        yield from self.line()

    def __attrs_post_init__(self) -> None:
        """You are in a twisty little maze of passages."""
        v = self.visit_stmt
        self.visit_if_stmt = partial(v, keywords={'if', 'else', 'elif'})
        self.visit_while_stmt = partial(v, keywords={'while', 'else'})
        self.visit_for_stmt = partial(v, keywords={'for', 'else'})
        self.visit_try_stmt = partial(v, keywords={'try', 'except', 'else', 'finally'})
        self.visit_except_clause = partial(v, keywords={'except'})
        self.visit_funcdef = partial(v, keywords={'def'})
        self.visit_with_stmt = partial(v, keywords={'with'})
        self.visit_classdef = partial(v, keywords={'class'})
        self.visit_async_funcdef = self.visit_async_stmt
        self.visit_decorated = self.visit_decorators


BRACKET = {token.LPAR: token.RPAR, token.LSQB: token.RSQB, token.LBRACE: token.RBRACE}
OPENING_BRACKETS = set(BRACKET.keys())
CLOSING_BRACKETS = set(BRACKET.values())
BRACKETS = OPENING_BRACKETS | CLOSING_BRACKETS


def whitespace(leaf: Leaf) -> str:
    """Return whitespace prefix if needed for the given `leaf`."""
    NO = ''
    SPACE = ' '
    DOUBLESPACE = '  '
    t = leaf.type
    p = leaf.parent
    v = leaf.value
    if t == token.COLON:
        return NO

    if t == token.COMMA:
        return NO

    if t == token.RPAR:
        return NO

    if t == token.COMMENT:
        return DOUBLESPACE

    if t == STANDALONE_COMMENT:
        return NO

    assert p is not None, f"INTERNAL ERROR: hand-made leaf without parent: {leaf!r}"
    if p.type in {syms.parameters, syms.arglist}:
        # untyped function signatures or calls
        if t == token.RPAR:
            return NO

        prev = leaf.prev_sibling
        if not prev or prev.type != token.COMMA:
            return NO

    if p.type == syms.varargslist:
        # lambdas
        if t == token.RPAR:
            return NO

        prev = leaf.prev_sibling
        if prev and prev.type != token.COMMA:
            return NO

    elif p.type == syms.typedargslist:
        # typed function signatures
        prev = leaf.prev_sibling
        if not prev:
            return NO

        if t == token.EQUAL:
            if prev.type != syms.tname:
                return NO

        elif prev.type == token.EQUAL:
            # A bit hacky: if the equal sign has whitespace, it means we
            # previously found it's a typed argument.  So, we're using that, too.
            return prev.prefix

        elif prev.type != token.COMMA:
            return NO

    elif p.type == syms.tname:
        # type names
        prev = leaf.prev_sibling
        if not prev:
            prevp = preceding_leaf(p)
            if not prevp or prevp.type != token.COMMA:
                return NO

    elif p.type == syms.trailer:
        # attributes and calls
        if t == token.LPAR or t == token.RPAR:
            return NO

        prev = leaf.prev_sibling
        if not prev:
            if t == token.DOT:
                prevp = preceding_leaf(p)
                if not prevp or prevp.type != token.NUMBER:
                    return NO

            elif t == token.LSQB:
                return NO

        elif prev.type != token.COMMA:
            return NO

    elif p.type == syms.argument:
        # single argument
        if t == token.EQUAL:
            return NO

        prev = leaf.prev_sibling
        if not prev:
            prevp = preceding_leaf(p)
            if not prevp or prevp.type == token.LPAR:
                return NO

        elif prev.type == token.EQUAL or prev.type == token.DOUBLESTAR:
            return NO

    elif p.type == syms.decorator:
        # decorators
        return NO

    elif p.type == syms.dotted_name:
        prev = leaf.prev_sibling
        if prev:
            return NO

        prevp = preceding_leaf(p)
        if not prevp or prevp.type == token.AT or prevp.type == token.DOT:
            return NO

    elif p.type == syms.classdef:
        if t == token.LPAR:
            return NO

        prev = leaf.prev_sibling
        if prev and prev.type == token.LPAR:
            return NO

    elif p.type == syms.subscript:
        # indexing
        if t == token.COLON:
            return NO

        prev = leaf.prev_sibling
        if not prev or prev.type == token.COLON:
            return NO

    elif p.type in {
        syms.test,
        syms.not_test,
        syms.xor_expr,
        syms.or_test,
        syms.and_test,
        syms.arith_expr,
        syms.shift_expr,
        syms.yield_expr,
        syms.term,
        syms.power,
        syms.comparison,
    }:
        # various arithmetic and logic expressions
        prev = leaf.prev_sibling
        if not prev:
            prevp = preceding_leaf(p)
            if not prevp or prevp.type in OPENING_BRACKETS:
                return NO

            if prevp.type == token.EQUAL:
                if prevp.parent and prevp.parent.type in {
                    syms.varargslist, syms.parameters, syms.arglist, syms.argument
                }:
                    return NO

        return SPACE

    elif p.type == syms.atom:
        if t in CLOSING_BRACKETS:
            return NO

        prev = leaf.prev_sibling
        if not prev:
            prevp = preceding_leaf(p)
            if not prevp:
                return NO

            if prevp.type in OPENING_BRACKETS:
                return NO

            if prevp.type == token.EQUAL:
                if prevp.parent and prevp.parent.type in {
                    syms.varargslist, syms.parameters, syms.arglist, syms.argument
                }:
                    return NO

            if prevp.type == token.DOUBLESTAR:
                if prevp.parent and prevp.parent.type in {
                    syms.varargslist, syms.parameters, syms.arglist, syms.dictsetmaker
                }:
                    return NO

        elif prev.type in OPENING_BRACKETS:
            return NO

        elif t == token.DOT:
            # dots, but not the first one.
            return NO

    elif (
        p.type == syms.listmaker or
        p.type == syms.testlist_gexp or
        p.type == syms.subscriptlist
    ):
        # list interior, including unpacking
        prev = leaf.prev_sibling
        if not prev:
            return NO

    elif p.type == syms.dictsetmaker:
        # dict and set interior, including unpacking
        prev = leaf.prev_sibling
        if not prev:
            return NO

        if prev.type == token.DOUBLESTAR:
            return NO

    elif p.type == syms.factor or p.type == syms.star_expr:
        # unary ops
        prev = leaf.prev_sibling
        if not prev:
            prevp = preceding_leaf(p)
            if not prevp or prevp.type in OPENING_BRACKETS:
                return NO

            prevp_parent = prevp.parent
            assert prevp_parent is not None
            if prevp.type == token.COLON and prevp_parent.type in {
                syms.subscript, syms.sliceop
            }:
                return NO

            elif prevp.type == token.EQUAL and prevp_parent.type == syms.argument:
                return NO

        elif t == token.NAME or t == token.NUMBER:
            return NO

    elif p.type == syms.import_from:
        if t == token.DOT:
            prev = leaf.prev_sibling
            if prev and prev.type == token.DOT:
                return NO

        elif t == token.NAME:
            if v == 'import':
                return SPACE

            prev = leaf.prev_sibling
            if prev and prev.type == token.DOT:
                return NO

    elif p.type == syms.sliceop:
        return NO

    return SPACE


def preceding_leaf(node: Optional[LN]) -> Optional[Leaf]:
    """Returns the first leaf that precedes `node`, if any."""
    while node:
        res = node.prev_sibling
        if res:
            if isinstance(res, Leaf):
                return res

            try:
                return list(res.leaves())[-1]

            except IndexError:
                return None

        node = node.parent
    return None


def is_delimiter(leaf: Leaf) -> int:
    """Returns the priority of the `leaf` delimiter. Returns 0 if not delimiter.

    Higher numbers are higher priority.
    """
    if leaf.type == token.COMMA:
        return COMMA_PRIORITY

    if leaf.type == token.NAME and leaf.value in LOGIC_OPERATORS:
        return LOGIC_PRIORITY

    if leaf.type in COMPARATORS:
        return COMPARATOR_PRIORITY

    if (
        leaf.type in MATH_OPERATORS and
        leaf.parent and
        leaf.parent.type not in {syms.factor, syms.star_expr}
    ):
        return MATH_PRIORITY

    return 0


def generate_comments(leaf: Leaf) -> Iterator[Leaf]:
    """Cleans the prefix of the `leaf` and generates comments from it, if any.

    Comments in lib2to3 are shoved into the whitespace prefix.  This happens
    in `pgen2/driver.py:Driver.parse_tokens()`.  This was a brilliant implementation
    move because it does away with modifying the grammar to include all the
    possible places in which comments can be placed.

    The sad consequence for us though is that comments don't "belong" anywhere.
    This is why this function generates simple parentless Leaf objects for
    comments.  We simply don't know what the correct parent should be.

    No matter though, we can live without this.  We really only need to
    differentiate between inline and standalone comments.  The latter don't
    share the line with any code.

    Inline comments are emitted as regular token.COMMENT leaves.  Standalone
    are emitted with a fake STANDALONE_COMMENT token identifier.
    """
    if not leaf.prefix:
        return

    if '#' not in leaf.prefix:
        return

    before_comment, content = leaf.prefix.split('#', 1)
    content = content.rstrip()
    if content and (content[0] not in {' ', '!', '#'}):
        content = ' ' + content
    is_standalone_comment = (
        '\n' in before_comment or '\n' in content or leaf.type == token.DEDENT
    )
    if not is_standalone_comment:
        # simple trailing comment
        yield Leaf(token.COMMENT, value='#' + content)
        return

    for line in ('#' + content).split('\n'):
        line = line.lstrip()
        if not line.startswith('#'):
            continue

        yield Leaf(STANDALONE_COMMENT, line)


def split_line(line: Line, line_length: int, inner: bool = False) -> Iterator[Line]:
    """Splits a `line` into potentially many lines.

    They should fit in the allotted `line_length` but might not be able to.
    `inner` signifies that there were a pair of brackets somewhere around the
    current `line`, possibly transitively. This means we can fallback to splitting
    by delimiters if the LHS/RHS don't yield any results.
    """
    line_str = str(line).strip('\n')
    if len(line_str) <= line_length and '\n' not in line_str:
        yield line
        return

    if line.is_def:
        split_funcs = [left_hand_split]
    elif line.inside_brackets:
        split_funcs = [delimiter_split]
        if '\n' not in line_str:
            # Only attempt RHS if we don't have multiline strings or comments
            # on this line.
            split_funcs.append(right_hand_split)
    else:
        split_funcs = [right_hand_split]
    for split_func in split_funcs:
        # We are accumulating lines in `result` because we might want to abort
        # mission and return the original line in the end, or attempt a different
        # split altogether.
        result: List[Line] = []
        try:
            for l in split_func(line):
                if str(l).strip('\n') == line_str:
                    raise CannotSplit("Split function returned an unchanged result")

                result.extend(split_line(l, line_length=line_length, inner=True))
        except CannotSplit as cs:
            continue

        else:
            yield from result
            break

    else:
        yield line


def left_hand_split(line: Line) -> Iterator[Line]:
    """Split line into many lines, starting with the first matching bracket pair.

    Note: this usually looks weird, only use this for function definitions.
    Prefer RHS otherwise.
    """
    head = Line(depth=line.depth)
    body = Line(depth=line.depth + 1, inside_brackets=True)
    tail = Line(depth=line.depth)
    tail_leaves: List[Leaf] = []
    body_leaves: List[Leaf] = []
    head_leaves: List[Leaf] = []
    current_leaves = head_leaves
    matching_bracket = None
    for leaf in line.leaves:
        if (
            current_leaves is body_leaves and
            leaf.type in CLOSING_BRACKETS and
            leaf.opening_bracket is matching_bracket  # type: ignore
        ):
            current_leaves = tail_leaves
        current_leaves.append(leaf)
        if current_leaves is head_leaves:
            if leaf.type in OPENING_BRACKETS:
                matching_bracket = leaf
                current_leaves = body_leaves
    # Since body is a new indent level, remove spurious leading whitespace.
    if body_leaves:
        normalize_prefix(body_leaves[0])
    # Build the new lines.
    for result, leaves in (
        (head, head_leaves), (body, body_leaves), (tail, tail_leaves)
    ):
        for leaf in leaves:
            result.append(leaf, preformatted=True)
            comment_after = line.comments.get(id(leaf))
            if comment_after:
                result.append(comment_after, preformatted=True)
    # Check if the split succeeded.
    tail_len = len(str(tail))
    if not body:
        if tail_len == 0:
            raise CannotSplit("Splitting brackets produced the same line")

        elif tail_len < 3:
            raise CannotSplit(
                f"Splitting brackets on an empty body to save "
                f"{tail_len} characters is not worth it"
            )

    for result in (head, body, tail):
        if result:
            yield result


def right_hand_split(line: Line) -> Iterator[Line]:
    """Split line into many lines, starting with the last matching bracket pair."""
    head = Line(depth=line.depth)
    body = Line(depth=line.depth + 1, inside_brackets=True)
    tail = Line(depth=line.depth)
    tail_leaves: List[Leaf] = []
    body_leaves: List[Leaf] = []
    head_leaves: List[Leaf] = []
    current_leaves = tail_leaves
    opening_bracket = None
    for leaf in reversed(line.leaves):
        if current_leaves is body_leaves:
            if leaf is opening_bracket:
                current_leaves = head_leaves
        current_leaves.append(leaf)
        if current_leaves is tail_leaves:
            if leaf.type in CLOSING_BRACKETS:
                opening_bracket = leaf.opening_bracket  # type: ignore
                current_leaves = body_leaves
    tail_leaves.reverse()
    body_leaves.reverse()
    head_leaves.reverse()
    # Since body is a new indent level, remove spurious leading whitespace.
    if body_leaves:
        normalize_prefix(body_leaves[0])
    # Build the new lines.
    for result, leaves in (
        (head, head_leaves), (body, body_leaves), (tail, tail_leaves)
    ):
        for leaf in leaves:
            result.append(leaf, preformatted=True)
            comment_after = line.comments.get(id(leaf))
            if comment_after:
                result.append(comment_after, preformatted=True)
    # Check if the split succeeded.
    tail_len = len(str(tail).strip('\n'))
    if not body:
        if tail_len == 0:
            raise CannotSplit("Splitting brackets produced the same line")

        elif tail_len < 3:
            raise CannotSplit(
                f"Splitting brackets on an empty body to save "
                f"{tail_len} characters is not worth it"
            )

    for result in (head, body, tail):
        if result:
            yield result


def delimiter_split(line: Line) -> Iterator[Line]:
    """Split according to delimiters of the highest priority.

    This kind of split doesn't increase indentation.
    """
    try:
        last_leaf = line.leaves[-1]
    except IndexError:
        raise CannotSplit("Line empty")

    delimiters = line.bracket_tracker.delimiters
    try:
        delimiter_priority = line.bracket_tracker.max_priority(exclude={id(last_leaf)})
    except ValueError:
        raise CannotSplit("No delimiters found")

    current_line = Line(depth=line.depth, inside_brackets=line.inside_brackets)
    for leaf in line.leaves:
        current_line.append(leaf, preformatted=True)
        comment_after = line.comments.get(id(leaf))
        if comment_after:
            current_line.append(comment_after, preformatted=True)
        leaf_priority = delimiters.get(id(leaf))
        if leaf_priority == delimiter_priority:
            normalize_prefix(current_line.leaves[0])
            yield current_line

            current_line = Line(depth=line.depth, inside_brackets=line.inside_brackets)
    if current_line:
        if (
            delimiter_priority == COMMA_PRIORITY and
            current_line.leaves[-1].type != token.COMMA
        ):
            current_line.append(Leaf(token.COMMA, ','))
        normalize_prefix(current_line.leaves[0])
        yield current_line


def is_import(leaf: Leaf) -> bool:
    """Returns True if the given leaf starts an import statement."""
    p = leaf.parent
    t = leaf.type
    v = leaf.value
    return bool(
        t == token.NAME and
        (
            (v == 'import' and p and p.type == syms.import_name) or
            (v == 'from' and p and p.type == syms.import_from)
        )
    )


def normalize_prefix(leaf: Leaf) -> None:
    """Leave existing extra newlines for imports.  Remove everything else."""
    if is_import(leaf):
        spl = leaf.prefix.split('#', 1)
        nl_count = spl[0].count('\n')
        if len(spl) > 1:
            # Skip one newline since it was for a standalone comment.
            nl_count -= 1
        leaf.prefix = '\n' * nl_count
        return

    leaf.prefix = ''


PYTHON_EXTENSIONS = {'.py'}
BLACKLISTED_DIRECTORIES = {
    'build', 'buck-out', 'dist', '_build', '.git', '.hg', '.mypy_cache', '.tox', '.venv'
}


def gen_python_files_in_dir(path: Path) -> Iterator[Path]:
    for child in path.iterdir():
        if child.is_dir():
            if child.name in BLACKLISTED_DIRECTORIES:
                continue

            yield from gen_python_files_in_dir(child)

        elif child.suffix in PYTHON_EXTENSIONS:
            yield child


@dataclass
class Report:
    """Provides a reformatting counter."""
    change_count: int = attrib(default=0)
    same_count: int = attrib(default=0)
    failure_count: int = attrib(default=0)

    def done(self, src: Path, changed: bool) -> None:
        """Increment the counter for successful reformatting. Write out a message."""
        if changed:
            out(f'reformatted {src}')
            self.change_count += 1
        else:
            out(f'{src} already well formatted, good job.', bold=False)
            self.same_count += 1

    def failed(self, src: Path, message: str) -> None:
        """Increment the counter for failed reformatting. Write out a message."""
        err(f'error: cannot format {src}: {message}')
        self.failure_count += 1

    @property
    def return_code(self) -> int:
        """Which return code should the app use considering the current state."""
        return 1 if self.failure_count else 0

    def __str__(self) -> str:
        """A color report of the current state.

        Use `click.unstyle` to remove colors.
        """
        report = []
        if self.change_count:
            s = 's' if self.change_count > 1 else ''
            report.append(
                click.style(f'{self.change_count} file{s} reformatted', bold=True)
            )
        if self.same_count:
            s = 's' if self.same_count > 1 else ''
            report.append(f'{self.same_count} file{s} left unchanged')
        if self.failure_count:
            s = 's' if self.failure_count > 1 else ''
            report.append(
                click.style(
                    f'{self.failure_count} file{s} failed to reformat', fg='red'
                )
            )
        return ', '.join(report) + '.'


def assert_equivalent(src: str, dst: str) -> None:
    """Raises AssertionError if `src` and `dst` aren't equivalent.

    This is a temporary sanity check until Black becomes stable.
    """

    import ast
    import traceback

    def _v(node: ast.AST, depth: int = 0) -> Iterator[str]:
        """Simple visitor generating strings to compare ASTs by content."""
        yield f"{'  ' * depth}{node.__class__.__name__}("

        for field in sorted(node._fields):
            try:
                value = getattr(node, field)
            except AttributeError:
                continue

            yield f"{'  ' * (depth+1)}{field}="

            if isinstance(value, list):
                for item in value:
                    if isinstance(item, ast.AST):
                        yield from _v(item, depth + 2)

            elif isinstance(value, ast.AST):
                yield from _v(value, depth + 2)

            else:
                yield f"{'  ' * (depth+2)}{value!r},  # {value.__class__.__name__}"

        yield f"{'  ' * depth})  # /{node.__class__.__name__}"

    try:
        src_ast = ast.parse(src)
    except Exception as exc:
        raise AssertionError(f"cannot parse source: {exc}") from None

    try:
        dst_ast = ast.parse(dst)
    except Exception as exc:
        log = dump_to_file(''.join(traceback.format_tb(exc.__traceback__)), dst)
        raise AssertionError(
            f"INTERNAL ERROR: Black produced invalid code: {exc}. "
            f"Please report a bug on https://github.com/ambv/black/issues.  "
            f"This invalid output might be helpful: {log}",
        ) from None

    src_ast_str = '\n'.join(_v(src_ast))
    dst_ast_str = '\n'.join(_v(dst_ast))
    if src_ast_str != dst_ast_str:
        log = dump_to_file(diff(src_ast_str, dst_ast_str, 'src', 'dst'))
        raise AssertionError(
            f"INTERNAL ERROR: Black produced code that is not equivalent to "
            f"the source.  "
            f"Please report a bug on https://github.com/ambv/black/issues.  "
            f"This diff might be helpful: {log}",
        ) from None


def assert_stable(src: str, dst: str, line_length: int) -> None:
    """Raises AssertionError if `dst` reformats differently the second time.

    This is a temporary sanity check until Black becomes stable.
    """
    newdst = format_str(dst, line_length=line_length)
    if dst != newdst:
        log = dump_to_file(
            diff(src, dst, 'source', 'first pass'),
            diff(dst, newdst, 'first pass', 'second pass'),
        )
        raise AssertionError(
            f"INTERNAL ERROR: Black produced different code on the second pass "
            f"of the formatter.  "
            f"Please report a bug on https://github.com/ambv/black/issues.  "
            f"This diff might be helpful: {log}",
        ) from None


def dump_to_file(*output: str) -> str:
    """Dumps `output` to a temporary file. Returns path to the file."""
    import tempfile

    with tempfile.NamedTemporaryFile(
        mode='w', prefix='blk_', suffix='.log', delete=False
    ) as f:
        for lines in output:
            f.write(lines)
            f.write('\n')
    return f.name


def diff(a: str, b: str, a_name: str, b_name: str) -> str:
    """Returns a udiff string between strings `a` and `b`."""
    import difflib

    a_lines = [line + '\n' for line in a.split('\n')]
    b_lines = [line + '\n' for line in b.split('\n')]
    return ''.join(
        difflib.unified_diff(a_lines, b_lines, fromfile=a_name, tofile=b_name, n=5)
    )


if __name__ == '__main__':
    main()