contrib/grey.py
author Gregory Szorc <gregory.szorc@gmail.com>
Sun, 06 Oct 2019 17:59:15 -0400
changeset 43104 74802979dd9d
parent 43083 7054fd370430
permissions -rw-r--r--
py3: define and use pycompat.itervalues() .itervalues() only exists on Python 2. Python 3's equivalent is .values(). But we don't want to blindly use .values() everywhere because on Python 2, it will create a list, which will have performance implications. This commit introduces pycompat.itervalues() which will call the appropriate method on the passed object. We update all callers of obj.itervalues() to pycompat.itervalues(obj) instead. With this commit, the only source tranforming remaining is for iteritems(). Victory is near... Differential Revision: https://phab.mercurial-scm.org/D7013
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
43083
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
     1
# no-check-code because 3rd party
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
     2
import ast
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
     3
import asyncio
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
     4
from concurrent.futures import Executor, ProcessPoolExecutor
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
     5
from contextlib import contextmanager
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
     6
from datetime import datetime
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
     7
from enum import Enum
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
     8
from functools import lru_cache, partial, wraps
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
     9
import io
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    10
import itertools
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    11
import logging
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    12
from multiprocessing import Manager, freeze_support
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    13
import os
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    14
from pathlib import Path
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    15
import pickle
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    16
import re
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    17
import signal
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    18
import sys
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    19
import tempfile
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    20
import tokenize
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    21
import traceback
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    22
from typing import (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    23
    Any,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    24
    Callable,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    25
    Collection,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    26
    Dict,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    27
    Generator,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    28
    Generic,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    29
    Iterable,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    30
    Iterator,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    31
    List,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    32
    Optional,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    33
    Pattern,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    34
    Sequence,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    35
    Set,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    36
    Tuple,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    37
    TypeVar,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    38
    Union,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    39
    cast,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    40
)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    41
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    42
from appdirs import user_cache_dir
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    43
from attr import dataclass, evolve, Factory
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    44
import click
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    45
import toml
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    46
from typed_ast import ast3, ast27
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    47
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    48
# lib2to3 fork
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    49
from blib2to3.pytree import Node, Leaf, type_repr
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    50
from blib2to3 import pygram, pytree
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    51
from blib2to3.pgen2 import driver, token
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    52
from blib2to3.pgen2.grammar import Grammar
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    53
from blib2to3.pgen2.parse import ParseError
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    54
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    55
__version__ = '19.3b1.dev95+gdc1add6.d20191005'
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    56
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    57
DEFAULT_LINE_LENGTH = 88
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    58
DEFAULT_EXCLUDES = (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    59
    r"/(\.eggs|\.git|\.hg|\.mypy_cache|\.nox|\.tox|\.venv|_build|buck-out|build|dist)/"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    60
)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    61
DEFAULT_INCLUDES = r"\.pyi?$"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    62
CACHE_DIR = Path(user_cache_dir("black", version=__version__))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    63
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    64
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    65
# types
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    66
FileContent = str
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    67
Encoding = str
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    68
NewLine = str
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    69
Depth = int
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    70
NodeType = int
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    71
LeafID = int
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    72
Priority = int
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    73
Index = int
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    74
LN = Union[Leaf, Node]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    75
SplitFunc = Callable[["Line", Collection["Feature"]], Iterator["Line"]]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    76
Timestamp = float
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    77
FileSize = int
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    78
CacheInfo = Tuple[Timestamp, FileSize]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    79
Cache = Dict[Path, CacheInfo]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    80
out = partial(click.secho, bold=True, err=True)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    81
err = partial(click.secho, fg="red", err=True)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    82
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    83
pygram.initialize(CACHE_DIR)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    84
syms = pygram.python_symbols
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    85
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    86
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    87
class NothingChanged(UserWarning):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    88
    """Raised when reformatted code is the same as source."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    89
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    90
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    91
class CannotSplit(Exception):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    92
    """A readable split that fits the allotted line length is impossible."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    93
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    94
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    95
class InvalidInput(ValueError):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    96
    """Raised when input source code fails all parse attempts."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    97
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    98
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
    99
class WriteBack(Enum):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   100
    NO = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   101
    YES = 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   102
    DIFF = 2
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   103
    CHECK = 3
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   104
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   105
    @classmethod
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   106
    def from_configuration(cls, *, check: bool, diff: bool) -> "WriteBack":
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   107
        if check and not diff:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   108
            return cls.CHECK
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   109
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   110
        return cls.DIFF if diff else cls.YES
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   111
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   112
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   113
class Changed(Enum):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   114
    NO = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   115
    CACHED = 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   116
    YES = 2
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   117
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   118
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   119
class TargetVersion(Enum):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   120
    PY27 = 2
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   121
    PY33 = 3
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   122
    PY34 = 4
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   123
    PY35 = 5
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   124
    PY36 = 6
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   125
    PY37 = 7
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   126
    PY38 = 8
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   127
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   128
    def is_python2(self) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   129
        return self is TargetVersion.PY27
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   130
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   131
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   132
PY36_VERSIONS = {TargetVersion.PY36, TargetVersion.PY37, TargetVersion.PY38}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   133
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   134
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   135
class Feature(Enum):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   136
    # All string literals are unicode
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   137
    UNICODE_LITERALS = 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   138
    F_STRINGS = 2
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   139
    NUMERIC_UNDERSCORES = 3
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   140
    TRAILING_COMMA_IN_CALL = 4
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   141
    TRAILING_COMMA_IN_DEF = 5
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   142
    # The following two feature-flags are mutually exclusive, and exactly one should be
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   143
    # set for every version of python.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   144
    ASYNC_IDENTIFIERS = 6
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   145
    ASYNC_KEYWORDS = 7
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   146
    ASSIGNMENT_EXPRESSIONS = 8
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   147
    POS_ONLY_ARGUMENTS = 9
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   148
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   149
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   150
VERSION_TO_FEATURES: Dict[TargetVersion, Set[Feature]] = {
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   151
    TargetVersion.PY27: {Feature.ASYNC_IDENTIFIERS},
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   152
    TargetVersion.PY33: {Feature.UNICODE_LITERALS, Feature.ASYNC_IDENTIFIERS},
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   153
    TargetVersion.PY34: {Feature.UNICODE_LITERALS, Feature.ASYNC_IDENTIFIERS},
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   154
    TargetVersion.PY35: {
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   155
        Feature.UNICODE_LITERALS,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   156
        Feature.TRAILING_COMMA_IN_CALL,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   157
        Feature.ASYNC_IDENTIFIERS,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   158
    },
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   159
    TargetVersion.PY36: {
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   160
        Feature.UNICODE_LITERALS,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   161
        Feature.F_STRINGS,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   162
        Feature.NUMERIC_UNDERSCORES,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   163
        Feature.TRAILING_COMMA_IN_CALL,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   164
        Feature.TRAILING_COMMA_IN_DEF,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   165
        Feature.ASYNC_IDENTIFIERS,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   166
    },
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   167
    TargetVersion.PY37: {
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   168
        Feature.UNICODE_LITERALS,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   169
        Feature.F_STRINGS,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   170
        Feature.NUMERIC_UNDERSCORES,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   171
        Feature.TRAILING_COMMA_IN_CALL,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   172
        Feature.TRAILING_COMMA_IN_DEF,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   173
        Feature.ASYNC_KEYWORDS,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   174
    },
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   175
    TargetVersion.PY38: {
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   176
        Feature.UNICODE_LITERALS,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   177
        Feature.F_STRINGS,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   178
        Feature.NUMERIC_UNDERSCORES,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   179
        Feature.TRAILING_COMMA_IN_CALL,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   180
        Feature.TRAILING_COMMA_IN_DEF,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   181
        Feature.ASYNC_KEYWORDS,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   182
        Feature.ASSIGNMENT_EXPRESSIONS,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   183
        Feature.POS_ONLY_ARGUMENTS,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   184
    },
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   185
}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   186
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   187
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   188
@dataclass
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   189
class FileMode:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   190
    target_versions: Set[TargetVersion] = Factory(set)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   191
    line_length: int = DEFAULT_LINE_LENGTH
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   192
    string_normalization: bool = True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   193
    is_pyi: bool = False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   194
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   195
    def get_cache_key(self) -> str:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   196
        if self.target_versions:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   197
            version_str = ",".join(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   198
                str(version.value)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   199
                for version in sorted(self.target_versions, key=lambda v: v.value)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   200
            )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   201
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   202
            version_str = "-"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   203
        parts = [
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   204
            version_str,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   205
            str(self.line_length),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   206
            str(int(self.string_normalization)),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   207
            str(int(self.is_pyi)),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   208
        ]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   209
        return ".".join(parts)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   210
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   211
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   212
def supports_feature(target_versions: Set[TargetVersion], feature: Feature) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   213
    return all(feature in VERSION_TO_FEATURES[version] for version in target_versions)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   214
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   215
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   216
def read_pyproject_toml(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   217
    ctx: click.Context, param: click.Parameter, value: Union[str, int, bool, None]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   218
) -> Optional[str]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   219
    """Inject Black configuration from "pyproject.toml" into defaults in `ctx`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   220
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   221
    Returns the path to a successfully found and read configuration file, None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   222
    otherwise.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   223
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   224
    assert not isinstance(value, (int, bool)), "Invalid parameter type passed"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   225
    if not value:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   226
        root = find_project_root(ctx.params.get("src", ()))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   227
        path = root / "pyproject.toml"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   228
        if path.is_file():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   229
            value = str(path)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   230
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   231
            return None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   232
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   233
    try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   234
        pyproject_toml = toml.load(value)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   235
        config = pyproject_toml.get("tool", {}).get("black", {})
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   236
    except (toml.TomlDecodeError, OSError) as e:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   237
        raise click.FileError(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   238
            filename=value, hint=f"Error reading configuration file: {e}"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   239
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   240
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   241
    if not config:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   242
        return None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   243
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   244
    if ctx.default_map is None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   245
        ctx.default_map = {}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   246
    ctx.default_map.update(  # type: ignore  # bad types in .pyi
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   247
        {k.replace("--", "").replace("-", "_"): v for k, v in config.items()}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   248
    )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   249
    return value
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   250
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   251
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   252
@click.command(context_settings=dict(help_option_names=["-h", "--help"]))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   253
@click.option("-c", "--code", type=str, help="Format the code passed in as a string.")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   254
@click.option(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   255
    "-l",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   256
    "--line-length",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   257
    type=int,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   258
    default=DEFAULT_LINE_LENGTH,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   259
    help="How many characters per line to allow.",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   260
    show_default=True,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   261
)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   262
@click.option(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   263
    "-t",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   264
    "--target-version",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   265
    type=click.Choice([v.name.lower() for v in TargetVersion]),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   266
    callback=lambda c, p, v: [TargetVersion[val.upper()] for val in v],
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   267
    multiple=True,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   268
    help=(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   269
        "Python versions that should be supported by Black's output. [default: "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   270
        "per-file auto-detection]"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   271
    ),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   272
)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   273
@click.option(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   274
    "--py36",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   275
    is_flag=True,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   276
    help=(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   277
        "Allow using Python 3.6-only syntax on all input files.  This will put "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   278
        "trailing commas in function signatures and calls also after *args and "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   279
        "**kwargs. Deprecated; use --target-version instead. "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   280
        "[default: per-file auto-detection]"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   281
    ),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   282
)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   283
@click.option(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   284
    "--pyi",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   285
    is_flag=True,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   286
    help=(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   287
        "Format all input files like typing stubs regardless of file extension "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   288
        "(useful when piping source on standard input)."
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   289
    ),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   290
)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   291
@click.option(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   292
    "-S",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   293
    "--skip-string-normalization",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   294
    is_flag=True,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   295
    help="Don't normalize string quotes or prefixes.",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   296
)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   297
@click.option(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   298
    "--check",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   299
    is_flag=True,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   300
    help=(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   301
        "Don't write the files back, just return the status.  Return code 0 "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   302
        "means nothing would change.  Return code 1 means some files would be "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   303
        "reformatted.  Return code 123 means there was an internal error."
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   304
    ),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   305
)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   306
@click.option(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   307
    "--diff",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   308
    is_flag=True,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   309
    help="Don't write the files back, just output a diff for each file on stdout.",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   310
)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   311
@click.option(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   312
    "--fast/--safe",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   313
    is_flag=True,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   314
    help="If --fast given, skip temporary sanity checks. [default: --safe]",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   315
)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   316
@click.option(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   317
    "--include",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   318
    type=str,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   319
    default=DEFAULT_INCLUDES,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   320
    help=(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   321
        "A regular expression that matches files and directories that should be "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   322
        "included on recursive searches.  An empty value means all files are "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   323
        "included regardless of the name.  Use forward slashes for directories on "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   324
        "all platforms (Windows, too).  Exclusions are calculated first, inclusions "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   325
        "later."
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   326
    ),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   327
    show_default=True,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   328
)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   329
@click.option(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   330
    "--exclude",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   331
    type=str,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   332
    default=DEFAULT_EXCLUDES,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   333
    help=(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   334
        "A regular expression that matches files and directories that should be "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   335
        "excluded on recursive searches.  An empty value means no paths are excluded. "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   336
        "Use forward slashes for directories on all platforms (Windows, too).  "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   337
        "Exclusions are calculated first, inclusions later."
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   338
    ),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   339
    show_default=True,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   340
)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   341
@click.option(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   342
    "-q",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   343
    "--quiet",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   344
    is_flag=True,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   345
    help=(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   346
        "Don't emit non-error messages to stderr. Errors are still emitted; "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   347
        "silence those with 2>/dev/null."
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   348
    ),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   349
)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   350
@click.option(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   351
    "-v",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   352
    "--verbose",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   353
    is_flag=True,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   354
    help=(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   355
        "Also emit messages to stderr about files that were not changed or were "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   356
        "ignored due to --exclude=."
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   357
    ),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   358
)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   359
@click.version_option(version=__version__)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   360
@click.argument(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   361
    "src",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   362
    nargs=-1,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   363
    type=click.Path(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   364
        exists=True, file_okay=True, dir_okay=True, readable=True, allow_dash=True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   365
    ),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   366
    is_eager=True,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   367
)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   368
@click.option(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   369
    "--config",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   370
    type=click.Path(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   371
        exists=False, file_okay=True, dir_okay=False, readable=True, allow_dash=False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   372
    ),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   373
    is_eager=True,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   374
    callback=read_pyproject_toml,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   375
    help="Read configuration from PATH.",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   376
)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   377
@click.pass_context
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   378
def main(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   379
    ctx: click.Context,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   380
    code: Optional[str],
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   381
    line_length: int,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   382
    target_version: List[TargetVersion],
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   383
    check: bool,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   384
    diff: bool,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   385
    fast: bool,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   386
    pyi: bool,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   387
    py36: bool,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   388
    skip_string_normalization: bool,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   389
    quiet: bool,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   390
    verbose: bool,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   391
    include: str,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   392
    exclude: str,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   393
    src: Tuple[str],
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   394
    config: Optional[str],
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   395
) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   396
    """The uncompromising code formatter."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   397
    write_back = WriteBack.from_configuration(check=check, diff=diff)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   398
    if target_version:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   399
        if py36:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   400
            err(f"Cannot use both --target-version and --py36")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   401
            ctx.exit(2)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   402
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   403
            versions = set(target_version)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   404
    elif py36:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   405
        err(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   406
            "--py36 is deprecated and will be removed in a future version. "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   407
            "Use --target-version py36 instead."
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   408
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   409
        versions = PY36_VERSIONS
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   410
    else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   411
        # We'll autodetect later.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   412
        versions = set()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   413
    mode = FileMode(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   414
        target_versions=versions,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   415
        line_length=line_length,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   416
        is_pyi=pyi,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   417
        string_normalization=not skip_string_normalization,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   418
    )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   419
    if config and verbose:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   420
        out(f"Using configuration from {config}.", bold=False, fg="blue")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   421
    if code is not None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   422
        print(format_str(code, mode=mode))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   423
        ctx.exit(0)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   424
    try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   425
        include_regex = re_compile_maybe_verbose(include)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   426
    except re.error:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   427
        err(f"Invalid regular expression for include given: {include!r}")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   428
        ctx.exit(2)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   429
    try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   430
        exclude_regex = re_compile_maybe_verbose(exclude)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   431
    except re.error:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   432
        err(f"Invalid regular expression for exclude given: {exclude!r}")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   433
        ctx.exit(2)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   434
    report = Report(check=check, quiet=quiet, verbose=verbose)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   435
    root = find_project_root(src)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   436
    sources: Set[Path] = set()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   437
    path_empty(src, quiet, verbose, ctx)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   438
    for s in src:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   439
        p = Path(s)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   440
        if p.is_dir():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   441
            sources.update(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   442
                gen_python_files_in_dir(p, root, include_regex, exclude_regex, report)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   443
            )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   444
        elif p.is_file() or s == "-":
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   445
            # if a file was explicitly given, we don't care about its extension
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   446
            sources.add(p)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   447
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   448
            err(f"invalid path: {s}")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   449
    if len(sources) == 0:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   450
        if verbose or not quiet:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   451
            out("No Python files are present to be formatted. Nothing to do 😴")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   452
        ctx.exit(0)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   453
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   454
    if len(sources) == 1:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   455
        reformat_one(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   456
            src=sources.pop(),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   457
            fast=fast,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   458
            write_back=write_back,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   459
            mode=mode,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   460
            report=report,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   461
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   462
    else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   463
        reformat_many(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   464
            sources=sources, fast=fast, write_back=write_back, mode=mode, report=report
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   465
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   466
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   467
    if verbose or not quiet:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   468
        out("Oh no! 💥 💔 💥" if report.return_code else "All done! ✨ 🍰 ✨")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   469
        click.secho(str(report), err=True)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   470
    ctx.exit(report.return_code)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   471
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   472
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   473
def path_empty(src: Tuple[str], quiet: bool, verbose: bool, ctx: click.Context) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   474
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   475
    Exit if there is no `src` provided for formatting
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   476
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   477
    if not src:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   478
        if verbose or not quiet:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   479
            out("No Path provided. Nothing to do 😴")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   480
            ctx.exit(0)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   481
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   482
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   483
def reformat_one(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   484
    src: Path, fast: bool, write_back: WriteBack, mode: FileMode, report: "Report"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   485
) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   486
    """Reformat a single file under `src` without spawning child processes.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   487
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   488
    `fast`, `write_back`, and `mode` options are passed to
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   489
    :func:`format_file_in_place` or :func:`format_stdin_to_stdout`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   490
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   491
    try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   492
        changed = Changed.NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   493
        if not src.is_file() and str(src) == "-":
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   494
            if format_stdin_to_stdout(fast=fast, write_back=write_back, mode=mode):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   495
                changed = Changed.YES
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   496
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   497
            cache: Cache = {}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   498
            if write_back != WriteBack.DIFF:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   499
                cache = read_cache(mode)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   500
                res_src = src.resolve()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   501
                if res_src in cache and cache[res_src] == get_cache_info(res_src):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   502
                    changed = Changed.CACHED
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   503
            if changed is not Changed.CACHED and format_file_in_place(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   504
                src, fast=fast, write_back=write_back, mode=mode
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   505
            ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   506
                changed = Changed.YES
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   507
            if (write_back is WriteBack.YES and changed is not Changed.CACHED) or (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   508
                write_back is WriteBack.CHECK and changed is Changed.NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   509
            ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   510
                write_cache(cache, [src], mode)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   511
        report.done(src, changed)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   512
    except Exception as exc:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   513
        report.failed(src, str(exc))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   514
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   515
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   516
def reformat_many(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   517
    sources: Set[Path],
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   518
    fast: bool,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   519
    write_back: WriteBack,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   520
    mode: FileMode,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   521
    report: "Report",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   522
) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   523
    """Reformat multiple files using a ProcessPoolExecutor."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   524
    loop = asyncio.get_event_loop()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   525
    worker_count = os.cpu_count()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   526
    if sys.platform == "win32":
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   527
        # Work around https://bugs.python.org/issue26903
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   528
        worker_count = min(worker_count, 61)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   529
    executor = ProcessPoolExecutor(max_workers=worker_count)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   530
    try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   531
        loop.run_until_complete(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   532
            schedule_formatting(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   533
                sources=sources,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   534
                fast=fast,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   535
                write_back=write_back,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   536
                mode=mode,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   537
                report=report,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   538
                loop=loop,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   539
                executor=executor,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   540
            )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   541
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   542
    finally:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   543
        shutdown(loop)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   544
        executor.shutdown()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   545
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   546
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   547
async def schedule_formatting(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   548
    sources: Set[Path],
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   549
    fast: bool,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   550
    write_back: WriteBack,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   551
    mode: FileMode,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   552
    report: "Report",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   553
    loop: asyncio.AbstractEventLoop,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   554
    executor: Executor,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   555
) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   556
    """Run formatting of `sources` in parallel using the provided `executor`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   557
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   558
    (Use ProcessPoolExecutors for actual parallelism.)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   559
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   560
    `write_back`, `fast`, and `mode` options are passed to
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   561
    :func:`format_file_in_place`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   562
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   563
    cache: Cache = {}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   564
    if write_back != WriteBack.DIFF:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   565
        cache = read_cache(mode)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   566
        sources, cached = filter_cached(cache, sources)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   567
        for src in sorted(cached):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   568
            report.done(src, Changed.CACHED)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   569
    if not sources:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   570
        return
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   571
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   572
    cancelled = []
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   573
    sources_to_cache = []
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   574
    lock = None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   575
    if write_back == WriteBack.DIFF:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   576
        # For diff output, we need locks to ensure we don't interleave output
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   577
        # from different processes.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   578
        manager = Manager()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   579
        lock = manager.Lock()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   580
    tasks = {
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   581
        asyncio.ensure_future(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   582
            loop.run_in_executor(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   583
                executor, format_file_in_place, src, fast, mode, write_back, lock
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   584
            )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   585
        ): src
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   586
        for src in sorted(sources)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   587
    }
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   588
    pending: Iterable[asyncio.Future] = tasks.keys()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   589
    try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   590
        loop.add_signal_handler(signal.SIGINT, cancel, pending)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   591
        loop.add_signal_handler(signal.SIGTERM, cancel, pending)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   592
    except NotImplementedError:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   593
        # There are no good alternatives for these on Windows.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   594
        pass
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   595
    while pending:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   596
        done, _ = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   597
        for task in done:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   598
            src = tasks.pop(task)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   599
            if task.cancelled():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   600
                cancelled.append(task)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   601
            elif task.exception():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   602
                report.failed(src, str(task.exception()))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   603
            else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   604
                changed = Changed.YES if task.result() else Changed.NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   605
                # If the file was written back or was successfully checked as
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   606
                # well-formatted, store this information in the cache.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   607
                if write_back is WriteBack.YES or (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   608
                    write_back is WriteBack.CHECK and changed is Changed.NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   609
                ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   610
                    sources_to_cache.append(src)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   611
                report.done(src, changed)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   612
    if cancelled:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   613
        await asyncio.gather(*cancelled, loop=loop, return_exceptions=True)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   614
    if sources_to_cache:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   615
        write_cache(cache, sources_to_cache, mode)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   616
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   617
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   618
def format_file_in_place(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   619
    src: Path,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   620
    fast: bool,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   621
    mode: FileMode,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   622
    write_back: WriteBack = WriteBack.NO,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   623
    lock: Any = None,  # multiprocessing.Manager().Lock() is some crazy proxy
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   624
) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   625
    """Format file under `src` path. Return True if changed.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   626
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   627
    If `write_back` is DIFF, write a diff to stdout. If it is YES, write reformatted
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   628
    code to the file.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   629
    `mode` and `fast` options are passed to :func:`format_file_contents`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   630
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   631
    if src.suffix == ".pyi":
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   632
        mode = evolve(mode, is_pyi=True)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   633
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   634
    then = datetime.utcfromtimestamp(src.stat().st_mtime)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   635
    with open(src, "rb") as buf:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   636
        src_contents, encoding, newline = decode_bytes(buf.read())
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   637
    try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   638
        dst_contents = format_file_contents(src_contents, fast=fast, mode=mode)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   639
    except NothingChanged:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   640
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   641
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   642
    if write_back == write_back.YES:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   643
        with open(src, "w", encoding=encoding, newline=newline) as f:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   644
            f.write(dst_contents)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   645
    elif write_back == write_back.DIFF:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   646
        now = datetime.utcnow()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   647
        src_name = f"{src}\t{then} +0000"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   648
        dst_name = f"{src}\t{now} +0000"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   649
        diff_contents = diff(src_contents, dst_contents, src_name, dst_name)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   650
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   651
        with lock or nullcontext():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   652
            f = io.TextIOWrapper(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   653
                sys.stdout.buffer,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   654
                encoding=encoding,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   655
                newline=newline,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   656
                write_through=True,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   657
            )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   658
            f.write(diff_contents)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   659
            f.detach()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   660
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   661
    return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   662
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   663
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   664
def format_stdin_to_stdout(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   665
    fast: bool, *, write_back: WriteBack = WriteBack.NO, mode: FileMode
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   666
) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   667
    """Format file on stdin. Return True if changed.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   668
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   669
    If `write_back` is YES, write reformatted code back to stdout. If it is DIFF,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   670
    write a diff to stdout. The `mode` argument is passed to
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   671
    :func:`format_file_contents`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   672
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   673
    then = datetime.utcnow()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   674
    src, encoding, newline = decode_bytes(sys.stdin.buffer.read())
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   675
    dst = src
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   676
    try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   677
        dst = format_file_contents(src, fast=fast, mode=mode)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   678
        return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   679
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   680
    except NothingChanged:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   681
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   682
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   683
    finally:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   684
        f = io.TextIOWrapper(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   685
            sys.stdout.buffer, encoding=encoding, newline=newline, write_through=True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   686
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   687
        if write_back == WriteBack.YES:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   688
            f.write(dst)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   689
        elif write_back == WriteBack.DIFF:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   690
            now = datetime.utcnow()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   691
            src_name = f"STDIN\t{then} +0000"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   692
            dst_name = f"STDOUT\t{now} +0000"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   693
            f.write(diff(src, dst, src_name, dst_name))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   694
        f.detach()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   695
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   696
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   697
def format_file_contents(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   698
    src_contents: str, *, fast: bool, mode: FileMode
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   699
) -> FileContent:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   700
    """Reformat contents a file and return new contents.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   701
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   702
    If `fast` is False, additionally confirm that the reformatted code is
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   703
    valid by calling :func:`assert_equivalent` and :func:`assert_stable` on it.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   704
    `mode` is passed to :func:`format_str`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   705
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   706
    if src_contents.strip() == "":
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   707
        raise NothingChanged
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   708
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   709
    dst_contents = format_str(src_contents, mode=mode)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   710
    if src_contents == dst_contents:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   711
        raise NothingChanged
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   712
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   713
    if not fast:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   714
        assert_equivalent(src_contents, dst_contents)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   715
        assert_stable(src_contents, dst_contents, mode=mode)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   716
    return dst_contents
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   717
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   718
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   719
def format_str(src_contents: str, *, mode: FileMode) -> FileContent:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   720
    """Reformat a string and return new contents.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   721
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   722
    `mode` determines formatting options, such as how many characters per line are
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   723
    allowed.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   724
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   725
    src_node = lib2to3_parse(src_contents.lstrip(), mode.target_versions)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   726
    dst_contents = []
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   727
    future_imports = get_future_imports(src_node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   728
    if mode.target_versions:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   729
        versions = mode.target_versions
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   730
    else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   731
        versions = detect_target_versions(src_node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   732
    normalize_fmt_off(src_node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   733
    lines = LineGenerator(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   734
        remove_u_prefix="unicode_literals" in future_imports
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   735
        or supports_feature(versions, Feature.UNICODE_LITERALS),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   736
        is_pyi=mode.is_pyi,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   737
        normalize_strings=mode.string_normalization,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   738
    )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   739
    elt = EmptyLineTracker(is_pyi=mode.is_pyi)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   740
    empty_line = Line()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   741
    after = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   742
    split_line_features = {
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   743
        feature
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   744
        for feature in {Feature.TRAILING_COMMA_IN_CALL, Feature.TRAILING_COMMA_IN_DEF}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   745
        if supports_feature(versions, feature)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   746
    }
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   747
    for current_line in lines.visit(src_node):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   748
        for _ in range(after):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   749
            dst_contents.append(str(empty_line))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   750
        before, after = elt.maybe_empty_lines(current_line)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   751
        for _ in range(before):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   752
            dst_contents.append(str(empty_line))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   753
        for line in split_line(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   754
            current_line, line_length=mode.line_length, features=split_line_features
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   755
        ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   756
            dst_contents.append(str(line))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   757
    return "".join(dst_contents)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   758
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   759
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   760
def decode_bytes(src: bytes) -> Tuple[FileContent, Encoding, NewLine]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   761
    """Return a tuple of (decoded_contents, encoding, newline).
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   762
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   763
    `newline` is either CRLF or LF but `decoded_contents` is decoded with
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   764
    universal newlines (i.e. only contains LF).
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   765
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   766
    srcbuf = io.BytesIO(src)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   767
    encoding, lines = tokenize.detect_encoding(srcbuf.readline)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   768
    if not lines:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   769
        return "", encoding, "\n"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   770
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   771
    newline = "\r\n" if b"\r\n" == lines[0][-2:] else "\n"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   772
    srcbuf.seek(0)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   773
    with io.TextIOWrapper(srcbuf, encoding) as tiow:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   774
        return tiow.read(), encoding, newline
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   775
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   776
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   777
def get_grammars(target_versions: Set[TargetVersion]) -> List[Grammar]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   778
    if not target_versions:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   779
        # No target_version specified, so try all grammars.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   780
        return [
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   781
            # Python 3.7+
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   782
            pygram.python_grammar_no_print_statement_no_exec_statement_async_keywords,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   783
            # Python 3.0-3.6
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   784
            pygram.python_grammar_no_print_statement_no_exec_statement,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   785
            # Python 2.7 with future print_function import
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   786
            pygram.python_grammar_no_print_statement,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   787
            # Python 2.7
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   788
            pygram.python_grammar,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   789
        ]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   790
    elif all(version.is_python2() for version in target_versions):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   791
        # Python 2-only code, so try Python 2 grammars.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   792
        return [
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   793
            # Python 2.7 with future print_function import
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   794
            pygram.python_grammar_no_print_statement,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   795
            # Python 2.7
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   796
            pygram.python_grammar,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   797
        ]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   798
    else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   799
        # Python 3-compatible code, so only try Python 3 grammar.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   800
        grammars = []
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   801
        # If we have to parse both, try to parse async as a keyword first
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   802
        if not supports_feature(target_versions, Feature.ASYNC_IDENTIFIERS):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   803
            # Python 3.7+
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   804
            grammars.append(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   805
                pygram.python_grammar_no_print_statement_no_exec_statement_async_keywords  # noqa: B950
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   806
            )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   807
        if not supports_feature(target_versions, Feature.ASYNC_KEYWORDS):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   808
            # Python 3.0-3.6
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   809
            grammars.append(pygram.python_grammar_no_print_statement_no_exec_statement)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   810
        # At least one of the above branches must have been taken, because every Python
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   811
        # version has exactly one of the two 'ASYNC_*' flags
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   812
        return grammars
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   813
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   814
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   815
def lib2to3_parse(src_txt: str, target_versions: Iterable[TargetVersion] = ()) -> Node:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   816
    """Given a string with source, return the lib2to3 Node."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   817
    if src_txt[-1:] != "\n":
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   818
        src_txt += "\n"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   819
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   820
    for grammar in get_grammars(set(target_versions)):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   821
        drv = driver.Driver(grammar, pytree.convert)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   822
        try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   823
            result = drv.parse_string(src_txt, True)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   824
            break
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   825
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   826
        except ParseError as pe:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   827
            lineno, column = pe.context[1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   828
            lines = src_txt.splitlines()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   829
            try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   830
                faulty_line = lines[lineno - 1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   831
            except IndexError:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   832
                faulty_line = "<line number missing in source>"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   833
            exc = InvalidInput(f"Cannot parse: {lineno}:{column}: {faulty_line}")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   834
    else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   835
        raise exc from None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   836
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   837
    if isinstance(result, Leaf):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   838
        result = Node(syms.file_input, [result])
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   839
    return result
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   840
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   841
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   842
def lib2to3_unparse(node: Node) -> str:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   843
    """Given a lib2to3 node, return its string representation."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   844
    code = str(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   845
    return code
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   846
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   847
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   848
T = TypeVar("T")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   849
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   850
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   851
class Visitor(Generic[T]):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   852
    """Basic lib2to3 visitor that yields things of type `T` on `visit()`."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   853
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   854
    def visit(self, node: LN) -> Iterator[T]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   855
        """Main method to visit `node` and its children.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   856
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   857
        It tries to find a `visit_*()` method for the given `node.type`, like
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   858
        `visit_simple_stmt` for Node objects or `visit_INDENT` for Leaf objects.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   859
        If no dedicated `visit_*()` method is found, chooses `visit_default()`
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   860
        instead.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   861
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   862
        Then yields objects of type `T` from the selected visitor.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   863
        """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   864
        if node.type < 256:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   865
            name = token.tok_name[node.type]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   866
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   867
            name = type_repr(node.type)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   868
        yield from getattr(self, f"visit_{name}", self.visit_default)(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   869
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   870
    def visit_default(self, node: LN) -> Iterator[T]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   871
        """Default `visit_*()` implementation. Recurses to children of `node`."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   872
        if isinstance(node, Node):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   873
            for child in node.children:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   874
                yield from self.visit(child)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   875
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   876
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   877
@dataclass
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   878
class DebugVisitor(Visitor[T]):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   879
    tree_depth: int = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   880
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   881
    def visit_default(self, node: LN) -> Iterator[T]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   882
        indent = " " * (2 * self.tree_depth)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   883
        if isinstance(node, Node):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   884
            _type = type_repr(node.type)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   885
            out(f"{indent}{_type}", fg="yellow")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   886
            self.tree_depth += 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   887
            for child in node.children:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   888
                yield from self.visit(child)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   889
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   890
            self.tree_depth -= 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   891
            out(f"{indent}/{_type}", fg="yellow", bold=False)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   892
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   893
            _type = token.tok_name.get(node.type, str(node.type))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   894
            out(f"{indent}{_type}", fg="blue", nl=False)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   895
            if node.prefix:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   896
                # We don't have to handle prefixes for `Node` objects since
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   897
                # that delegates to the first child anyway.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   898
                out(f" {node.prefix!r}", fg="green", bold=False, nl=False)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   899
            out(f" {node.value!r}", fg="blue", bold=False)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   900
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   901
    @classmethod
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   902
    def show(cls, code: Union[str, Leaf, Node]) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   903
        """Pretty-print the lib2to3 AST of a given string of `code`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   904
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   905
        Convenience method for debugging.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   906
        """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   907
        v: DebugVisitor[None] = DebugVisitor()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   908
        if isinstance(code, str):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   909
            code = lib2to3_parse(code)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   910
        list(v.visit(code))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   911
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   912
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   913
WHITESPACE = {token.DEDENT, token.INDENT, token.NEWLINE}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   914
STATEMENT = {
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   915
    syms.if_stmt,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   916
    syms.while_stmt,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   917
    syms.for_stmt,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   918
    syms.try_stmt,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   919
    syms.except_clause,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   920
    syms.with_stmt,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   921
    syms.funcdef,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   922
    syms.classdef,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   923
}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   924
STANDALONE_COMMENT = 153
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   925
token.tok_name[STANDALONE_COMMENT] = "STANDALONE_COMMENT"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   926
LOGIC_OPERATORS = {"and", "or"}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   927
COMPARATORS = {
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   928
    token.LESS,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   929
    token.GREATER,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   930
    token.EQEQUAL,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   931
    token.NOTEQUAL,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   932
    token.LESSEQUAL,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   933
    token.GREATEREQUAL,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   934
}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   935
MATH_OPERATORS = {
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   936
    token.VBAR,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   937
    token.CIRCUMFLEX,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   938
    token.AMPER,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   939
    token.LEFTSHIFT,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   940
    token.RIGHTSHIFT,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   941
    token.PLUS,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   942
    token.MINUS,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   943
    token.STAR,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   944
    token.SLASH,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   945
    token.DOUBLESLASH,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   946
    token.PERCENT,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   947
    token.AT,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   948
    token.TILDE,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   949
    token.DOUBLESTAR,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   950
}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   951
STARS = {token.STAR, token.DOUBLESTAR}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   952
VARARGS_SPECIALS = STARS | {token.SLASH}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   953
VARARGS_PARENTS = {
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   954
    syms.arglist,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   955
    syms.argument,  # double star in arglist
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   956
    syms.trailer,  # single argument to call
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   957
    syms.typedargslist,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   958
    syms.varargslist,  # lambdas
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   959
}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   960
UNPACKING_PARENTS = {
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   961
    syms.atom,  # single element of a list or set literal
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   962
    syms.dictsetmaker,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   963
    syms.listmaker,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   964
    syms.testlist_gexp,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   965
    syms.testlist_star_expr,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   966
}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   967
TEST_DESCENDANTS = {
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   968
    syms.test,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   969
    syms.lambdef,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   970
    syms.or_test,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   971
    syms.and_test,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   972
    syms.not_test,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   973
    syms.comparison,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   974
    syms.star_expr,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   975
    syms.expr,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   976
    syms.xor_expr,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   977
    syms.and_expr,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   978
    syms.shift_expr,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   979
    syms.arith_expr,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   980
    syms.trailer,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   981
    syms.term,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   982
    syms.power,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   983
}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   984
ASSIGNMENTS = {
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   985
    "=",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   986
    "+=",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   987
    "-=",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   988
    "*=",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   989
    "@=",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   990
    "/=",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   991
    "%=",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   992
    "&=",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   993
    "|=",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   994
    "^=",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   995
    "<<=",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   996
    ">>=",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   997
    "**=",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   998
    "//=",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
   999
}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1000
COMPREHENSION_PRIORITY = 20
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1001
COMMA_PRIORITY = 18
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1002
TERNARY_PRIORITY = 16
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1003
LOGIC_PRIORITY = 14
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1004
STRING_PRIORITY = 12
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1005
COMPARATOR_PRIORITY = 10
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1006
MATH_PRIORITIES = {
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1007
    token.VBAR: 9,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1008
    token.CIRCUMFLEX: 8,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1009
    token.AMPER: 7,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1010
    token.LEFTSHIFT: 6,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1011
    token.RIGHTSHIFT: 6,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1012
    token.PLUS: 5,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1013
    token.MINUS: 5,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1014
    token.STAR: 4,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1015
    token.SLASH: 4,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1016
    token.DOUBLESLASH: 4,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1017
    token.PERCENT: 4,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1018
    token.AT: 4,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1019
    token.TILDE: 3,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1020
    token.DOUBLESTAR: 2,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1021
}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1022
DOT_PRIORITY = 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1023
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1024
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1025
@dataclass
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1026
class BracketTracker:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1027
    """Keeps track of brackets on a line."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1028
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1029
    depth: int = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1030
    bracket_match: Dict[Tuple[Depth, NodeType], Leaf] = Factory(dict)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1031
    delimiters: Dict[LeafID, Priority] = Factory(dict)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1032
    previous: Optional[Leaf] = None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1033
    _for_loop_depths: List[int] = Factory(list)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1034
    _lambda_argument_depths: List[int] = Factory(list)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1035
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1036
    def mark(self, leaf: Leaf) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1037
        """Mark `leaf` with bracket-related metadata. Keep track of delimiters.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1038
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1039
        All leaves receive an int `bracket_depth` field that stores how deep
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1040
        within brackets a given leaf is. 0 means there are no enclosing brackets
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1041
        that started on this line.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1042
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1043
        If a leaf is itself a closing bracket, it receives an `opening_bracket`
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1044
        field that it forms a pair with. This is a one-directional link to
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1045
        avoid reference cycles.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1046
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1047
        If a leaf is a delimiter (a token on which Black can split the line if
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1048
        needed) and it's on depth 0, its `id()` is stored in the tracker's
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1049
        `delimiters` field.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1050
        """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1051
        if leaf.type == token.COMMENT:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1052
            return
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1053
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1054
        self.maybe_decrement_after_for_loop_variable(leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1055
        self.maybe_decrement_after_lambda_arguments(leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1056
        if leaf.type in CLOSING_BRACKETS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1057
            self.depth -= 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1058
            opening_bracket = self.bracket_match.pop((self.depth, leaf.type))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1059
            leaf.opening_bracket = opening_bracket
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1060
        leaf.bracket_depth = self.depth
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1061
        if self.depth == 0:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1062
            delim = is_split_before_delimiter(leaf, self.previous)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1063
            if delim and self.previous is not None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1064
                self.delimiters[id(self.previous)] = delim
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1065
            else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1066
                delim = is_split_after_delimiter(leaf, self.previous)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1067
                if delim:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1068
                    self.delimiters[id(leaf)] = delim
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1069
        if leaf.type in OPENING_BRACKETS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1070
            self.bracket_match[self.depth, BRACKET[leaf.type]] = leaf
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1071
            self.depth += 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1072
        self.previous = leaf
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1073
        self.maybe_increment_lambda_arguments(leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1074
        self.maybe_increment_for_loop_variable(leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1075
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1076
    def any_open_brackets(self) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1077
        """Return True if there is an yet unmatched open bracket on the line."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1078
        return bool(self.bracket_match)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1079
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1080
    def max_delimiter_priority(self, exclude: Iterable[LeafID] = ()) -> Priority:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1081
        """Return the highest priority of a delimiter found on the line.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1082
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1083
        Values are consistent with what `is_split_*_delimiter()` return.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1084
        Raises ValueError on no delimiters.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1085
        """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1086
        return max(v for k, v in self.delimiters.items() if k not in exclude)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1087
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1088
    def delimiter_count_with_priority(self, priority: Priority = 0) -> int:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1089
        """Return the number of delimiters with the given `priority`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1090
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1091
        If no `priority` is passed, defaults to max priority on the line.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1092
        """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1093
        if not self.delimiters:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1094
            return 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1095
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1096
        priority = priority or self.max_delimiter_priority()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1097
        return sum(1 for p in self.delimiters.values() if p == priority)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1098
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1099
    def maybe_increment_for_loop_variable(self, leaf: Leaf) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1100
        """In a for loop, or comprehension, the variables are often unpacks.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1101
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1102
        To avoid splitting on the comma in this situation, increase the depth of
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1103
        tokens between `for` and `in`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1104
        """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1105
        if leaf.type == token.NAME and leaf.value == "for":
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1106
            self.depth += 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1107
            self._for_loop_depths.append(self.depth)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1108
            return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1109
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1110
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1111
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1112
    def maybe_decrement_after_for_loop_variable(self, leaf: Leaf) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1113
        """See `maybe_increment_for_loop_variable` above for explanation."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1114
        if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1115
            self._for_loop_depths
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1116
            and self._for_loop_depths[-1] == self.depth
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1117
            and leaf.type == token.NAME
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1118
            and leaf.value == "in"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1119
        ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1120
            self.depth -= 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1121
            self._for_loop_depths.pop()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1122
            return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1123
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1124
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1125
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1126
    def maybe_increment_lambda_arguments(self, leaf: Leaf) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1127
        """In a lambda expression, there might be more than one argument.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1128
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1129
        To avoid splitting on the comma in this situation, increase the depth of
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1130
        tokens between `lambda` and `:`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1131
        """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1132
        if leaf.type == token.NAME and leaf.value == "lambda":
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1133
            self.depth += 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1134
            self._lambda_argument_depths.append(self.depth)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1135
            return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1136
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1137
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1138
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1139
    def maybe_decrement_after_lambda_arguments(self, leaf: Leaf) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1140
        """See `maybe_increment_lambda_arguments` above for explanation."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1141
        if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1142
            self._lambda_argument_depths
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1143
            and self._lambda_argument_depths[-1] == self.depth
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1144
            and leaf.type == token.COLON
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1145
        ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1146
            self.depth -= 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1147
            self._lambda_argument_depths.pop()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1148
            return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1149
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1150
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1151
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1152
    def get_open_lsqb(self) -> Optional[Leaf]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1153
        """Return the most recent opening square bracket (if any)."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1154
        return self.bracket_match.get((self.depth - 1, token.RSQB))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1155
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1156
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1157
@dataclass
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1158
class Line:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1159
    """Holds leaves and comments. Can be printed with `str(line)`."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1160
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1161
    depth: int = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1162
    leaves: List[Leaf] = Factory(list)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1163
    comments: Dict[LeafID, List[Leaf]] = Factory(dict)  # keys ordered like `leaves`
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1164
    bracket_tracker: BracketTracker = Factory(BracketTracker)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1165
    inside_brackets: bool = False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1166
    should_explode: bool = False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1167
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1168
    def append(self, leaf: Leaf, preformatted: bool = False) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1169
        """Add a new `leaf` to the end of the line.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1170
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1171
        Unless `preformatted` is True, the `leaf` will receive a new consistent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1172
        whitespace prefix and metadata applied by :class:`BracketTracker`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1173
        Trailing commas are maybe removed, unpacked for loop variables are
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1174
        demoted from being delimiters.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1175
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1176
        Inline comments are put aside.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1177
        """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1178
        has_value = leaf.type in BRACKETS or bool(leaf.value.strip())
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1179
        if not has_value:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1180
            return
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1181
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1182
        if token.COLON == leaf.type and self.is_class_paren_empty:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1183
            del self.leaves[-2:]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1184
        if self.leaves and not preformatted:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1185
            # Note: at this point leaf.prefix should be empty except for
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1186
            # imports, for which we only preserve newlines.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1187
            leaf.prefix += whitespace(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1188
                leaf, complex_subscript=self.is_complex_subscript(leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1189
            )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1190
        if self.inside_brackets or not preformatted:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1191
            self.bracket_tracker.mark(leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1192
            self.maybe_remove_trailing_comma(leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1193
        if not self.append_comment(leaf):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1194
            self.leaves.append(leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1195
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1196
    def append_safe(self, leaf: Leaf, preformatted: bool = False) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1197
        """Like :func:`append()` but disallow invalid standalone comment structure.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1198
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1199
        Raises ValueError when any `leaf` is appended after a standalone comment
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1200
        or when a standalone comment is not the first leaf on the line.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1201
        """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1202
        if self.bracket_tracker.depth == 0:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1203
            if self.is_comment:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1204
                raise ValueError("cannot append to standalone comments")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1205
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1206
            if self.leaves and leaf.type == STANDALONE_COMMENT:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1207
                raise ValueError(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1208
                    "cannot append standalone comments to a populated line"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1209
                )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1210
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1211
        self.append(leaf, preformatted=preformatted)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1212
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1213
    @property
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1214
    def is_comment(self) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1215
        """Is this line a standalone comment?"""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1216
        return len(self.leaves) == 1 and self.leaves[0].type == STANDALONE_COMMENT
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1217
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1218
    @property
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1219
    def is_decorator(self) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1220
        """Is this line a decorator?"""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1221
        return bool(self) and self.leaves[0].type == token.AT
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1222
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1223
    @property
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1224
    def is_import(self) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1225
        """Is this an import line?"""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1226
        return bool(self) and is_import(self.leaves[0])
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1227
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1228
    @property
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1229
    def is_class(self) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1230
        """Is this line a class definition?"""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1231
        return (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1232
            bool(self)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1233
            and self.leaves[0].type == token.NAME
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1234
            and self.leaves[0].value == "class"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1235
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1236
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1237
    @property
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1238
    def is_stub_class(self) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1239
        """Is this line a class definition with a body consisting only of "..."?"""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1240
        return self.is_class and self.leaves[-3:] == [
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1241
            Leaf(token.DOT, ".") for _ in range(3)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1242
        ]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1243
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1244
    @property
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1245
    def is_collection_with_optional_trailing_comma(self) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1246
        """Is this line a collection literal with a trailing comma that's optional?
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1247
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1248
        Note that the trailing comma in a 1-tuple is not optional.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1249
        """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1250
        if not self.leaves or len(self.leaves) < 4:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1251
            return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1252
        # Look for and address a trailing colon.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1253
        if self.leaves[-1].type == token.COLON:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1254
            closer = self.leaves[-2]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1255
            close_index = -2
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1256
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1257
            closer = self.leaves[-1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1258
            close_index = -1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1259
        if closer.type not in CLOSING_BRACKETS or self.inside_brackets:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1260
            return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1261
        if closer.type == token.RPAR:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1262
            # Tuples require an extra check, because if there's only
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1263
            # one element in the tuple removing the comma unmakes the
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1264
            # tuple.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1265
            #
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1266
            # We also check for parens before looking for the trailing
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1267
            # comma because in some cases (eg assigning a dict
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1268
            # literal) the literal gets wrapped in temporary parens
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1269
            # during parsing. This case is covered by the
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1270
            # collections.py test data.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1271
            opener = closer.opening_bracket
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1272
            for _open_index, leaf in enumerate(self.leaves):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1273
                if leaf is opener:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1274
                    break
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1275
            else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1276
                # Couldn't find the matching opening paren, play it safe.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1277
                return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1278
            commas = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1279
            comma_depth = self.leaves[close_index - 1].bracket_depth
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1280
            for leaf in self.leaves[_open_index + 1 : close_index]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1281
                if leaf.bracket_depth == comma_depth and leaf.type == token.COMMA:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1282
                    commas += 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1283
            if commas > 1:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1284
                # We haven't looked yet for the trailing comma because
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1285
                # we might also have caught noop parens.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1286
                return self.leaves[close_index - 1].type == token.COMMA
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1287
            elif commas == 1:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1288
                return False  # it's either a one-tuple or didn't have a trailing comma
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1289
            if self.leaves[close_index - 1].type in CLOSING_BRACKETS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1290
                close_index -= 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1291
                closer = self.leaves[close_index]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1292
                if closer.type == token.RPAR:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1293
                    # TODO: this is a gut feeling. Will we ever see this?
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1294
                    return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1295
        if self.leaves[close_index - 1].type != token.COMMA:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1296
            return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1297
        return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1298
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1299
    @property
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1300
    def is_def(self) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1301
        """Is this a function definition? (Also returns True for async defs.)"""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1302
        try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1303
            first_leaf = self.leaves[0]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1304
        except IndexError:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1305
            return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1306
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1307
        try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1308
            second_leaf: Optional[Leaf] = self.leaves[1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1309
        except IndexError:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1310
            second_leaf = None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1311
        return (first_leaf.type == token.NAME and first_leaf.value == "def") or (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1312
            first_leaf.type == token.ASYNC
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1313
            and second_leaf is not None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1314
            and second_leaf.type == token.NAME
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1315
            and second_leaf.value == "def"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1316
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1317
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1318
    @property
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1319
    def is_class_paren_empty(self) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1320
        """Is this a class with no base classes but using parentheses?
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1321
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1322
        Those are unnecessary and should be removed.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1323
        """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1324
        return (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1325
            bool(self)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1326
            and len(self.leaves) == 4
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1327
            and self.is_class
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1328
            and self.leaves[2].type == token.LPAR
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1329
            and self.leaves[2].value == "("
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1330
            and self.leaves[3].type == token.RPAR
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1331
            and self.leaves[3].value == ")"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1332
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1333
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1334
    @property
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1335
    def is_triple_quoted_string(self) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1336
        """Is the line a triple quoted string?"""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1337
        return (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1338
            bool(self)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1339
            and self.leaves[0].type == token.STRING
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1340
            and self.leaves[0].value.startswith(('"""', "'''"))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1341
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1342
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1343
    def contains_standalone_comments(self, depth_limit: int = sys.maxsize) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1344
        """If so, needs to be split before emitting."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1345
        for leaf in self.leaves:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1346
            if leaf.type == STANDALONE_COMMENT:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1347
                if leaf.bracket_depth <= depth_limit:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1348
                    return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1349
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1350
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1351
    def contains_uncollapsable_type_comments(self) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1352
        ignored_ids = set()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1353
        try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1354
            last_leaf = self.leaves[-1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1355
            ignored_ids.add(id(last_leaf))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1356
            if last_leaf.type == token.COMMA or (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1357
                last_leaf.type == token.RPAR and not last_leaf.value
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1358
            ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1359
                # When trailing commas or optional parens are inserted by Black for
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1360
                # consistency, comments after the previous last element are not moved
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1361
                # (they don't have to, rendering will still be correct).  So we ignore
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1362
                # trailing commas and invisible.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1363
                last_leaf = self.leaves[-2]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1364
                ignored_ids.add(id(last_leaf))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1365
        except IndexError:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1366
            return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1367
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1368
        # A type comment is uncollapsable if it is attached to a leaf
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1369
        # that isn't at the end of the line (since that could cause it
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1370
        # to get associated to a different argument) or if there are
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1371
        # comments before it (since that could cause it to get hidden
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1372
        # behind a comment.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1373
        comment_seen = False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1374
        for leaf_id, comments in self.comments.items():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1375
            for comment in comments:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1376
                if is_type_comment(comment):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1377
                    if leaf_id not in ignored_ids or comment_seen:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1378
                        return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1379
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1380
            comment_seen = True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1381
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1382
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1383
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1384
    def contains_unsplittable_type_ignore(self) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1385
        if not self.leaves:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1386
            return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1387
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1388
        # If a 'type: ignore' is attached to the end of a line, we
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1389
        # can't split the line, because we can't know which of the
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1390
        # subexpressions the ignore was meant to apply to.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1391
        #
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1392
        # We only want this to apply to actual physical lines from the
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1393
        # original source, though: we don't want the presence of a
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1394
        # 'type: ignore' at the end of a multiline expression to
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1395
        # justify pushing it all onto one line. Thus we
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1396
        # (unfortunately) need to check the actual source lines and
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1397
        # only report an unsplittable 'type: ignore' if this line was
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1398
        # one line in the original code.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1399
        if self.leaves[0].lineno == self.leaves[-1].lineno:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1400
            for comment in self.comments.get(id(self.leaves[-1]), []):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1401
                if is_type_comment(comment, " ignore"):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1402
                    return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1403
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1404
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1405
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1406
    def contains_multiline_strings(self) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1407
        for leaf in self.leaves:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1408
            if is_multiline_string(leaf):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1409
                return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1410
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1411
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1412
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1413
    def maybe_remove_trailing_comma(self, closing: Leaf) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1414
        """Remove trailing comma if there is one and it's safe."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1415
        if not (self.leaves and self.leaves[-1].type == token.COMMA):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1416
            return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1417
        # We remove trailing commas only in the case of importing a
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1418
        # single name from a module.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1419
        if not (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1420
            self.leaves
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1421
            and self.is_import
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1422
            and len(self.leaves) > 4
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1423
            and self.leaves[-1].type == token.COMMA
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1424
            and closing.type in CLOSING_BRACKETS
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1425
            and self.leaves[-4].type == token.NAME
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1426
            and (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1427
                # regular `from foo import bar,`
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1428
                self.leaves[-4].value == "import"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1429
                # `from foo import (bar as baz,)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1430
                or (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1431
                    len(self.leaves) > 6
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1432
                    and self.leaves[-6].value == "import"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1433
                    and self.leaves[-3].value == "as"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1434
                )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1435
                # `from foo import bar as baz,`
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1436
                or (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1437
                    len(self.leaves) > 5
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1438
                    and self.leaves[-5].value == "import"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1439
                    and self.leaves[-3].value == "as"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1440
                )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1441
            )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1442
            and closing.type == token.RPAR
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1443
        ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1444
            return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1445
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1446
        self.remove_trailing_comma()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1447
        return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1448
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1449
    def append_comment(self, comment: Leaf) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1450
        """Add an inline or standalone comment to the line."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1451
        if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1452
            comment.type == STANDALONE_COMMENT
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1453
            and self.bracket_tracker.any_open_brackets()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1454
        ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1455
            comment.prefix = ""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1456
            return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1457
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1458
        if comment.type != token.COMMENT:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1459
            return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1460
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1461
        if not self.leaves:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1462
            comment.type = STANDALONE_COMMENT
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1463
            comment.prefix = ""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1464
            return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1465
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1466
        last_leaf = self.leaves[-1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1467
        if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1468
            last_leaf.type == token.RPAR
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1469
            and not last_leaf.value
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1470
            and last_leaf.parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1471
            and len(list(last_leaf.parent.leaves())) <= 3
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1472
            and not is_type_comment(comment)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1473
        ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1474
            # Comments on an optional parens wrapping a single leaf should belong to
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1475
            # the wrapped node except if it's a type comment. Pinning the comment like
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1476
            # this avoids unstable formatting caused by comment migration.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1477
            if len(self.leaves) < 2:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1478
                comment.type = STANDALONE_COMMENT
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1479
                comment.prefix = ""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1480
                return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1481
            last_leaf = self.leaves[-2]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1482
        self.comments.setdefault(id(last_leaf), []).append(comment)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1483
        return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1484
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1485
    def comments_after(self, leaf: Leaf) -> List[Leaf]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1486
        """Generate comments that should appear directly after `leaf`."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1487
        return self.comments.get(id(leaf), [])
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1488
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1489
    def remove_trailing_comma(self) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1490
        """Remove the trailing comma and moves the comments attached to it."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1491
        trailing_comma = self.leaves.pop()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1492
        trailing_comma_comments = self.comments.pop(id(trailing_comma), [])
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1493
        self.comments.setdefault(id(self.leaves[-1]), []).extend(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1494
            trailing_comma_comments
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1495
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1496
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1497
    def is_complex_subscript(self, leaf: Leaf) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1498
        """Return True iff `leaf` is part of a slice with non-trivial exprs."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1499
        open_lsqb = self.bracket_tracker.get_open_lsqb()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1500
        if open_lsqb is None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1501
            return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1502
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1503
        subscript_start = open_lsqb.next_sibling
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1504
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1505
        if isinstance(subscript_start, Node):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1506
            if subscript_start.type == syms.listmaker:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1507
                return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1508
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1509
            if subscript_start.type == syms.subscriptlist:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1510
                subscript_start = child_towards(subscript_start, leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1511
        return subscript_start is not None and any(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1512
            n.type in TEST_DESCENDANTS for n in subscript_start.pre_order()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1513
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1514
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1515
    def __str__(self) -> str:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1516
        """Render the line."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1517
        if not self:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1518
            return "\n"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1519
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1520
        indent = "    " * self.depth
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1521
        leaves = iter(self.leaves)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1522
        first = next(leaves)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1523
        res = f"{first.prefix}{indent}{first.value}"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1524
        for leaf in leaves:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1525
            res += str(leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1526
        for comment in itertools.chain.from_iterable(self.comments.values()):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1527
            res += str(comment)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1528
        return res + "\n"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1529
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1530
    def __bool__(self) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1531
        """Return True if the line has leaves or comments."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1532
        return bool(self.leaves or self.comments)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1533
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1534
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1535
@dataclass
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1536
class EmptyLineTracker:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1537
    """Provides a stateful method that returns the number of potential extra
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1538
    empty lines needed before and after the currently processed line.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1539
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1540
    Note: this tracker works on lines that haven't been split yet.  It assumes
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1541
    the prefix of the first leaf consists of optional newlines.  Those newlines
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1542
    are consumed by `maybe_empty_lines()` and included in the computation.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1543
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1544
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1545
    is_pyi: bool = False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1546
    previous_line: Optional[Line] = None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1547
    previous_after: int = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1548
    previous_defs: List[int] = Factory(list)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1549
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1550
    def maybe_empty_lines(self, current_line: Line) -> Tuple[int, int]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1551
        """Return the number of extra empty lines before and after the `current_line`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1552
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1553
        This is for separating `def`, `async def` and `class` with extra empty
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1554
        lines (two on module-level).
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1555
        """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1556
        before, after = self._maybe_empty_lines(current_line)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1557
        before = (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1558
            # Black should not insert empty lines at the beginning
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1559
            # of the file
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1560
            0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1561
            if self.previous_line is None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1562
            else before - self.previous_after
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1563
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1564
        self.previous_after = after
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1565
        self.previous_line = current_line
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1566
        return before, after
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1567
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1568
    def _maybe_empty_lines(self, current_line: Line) -> Tuple[int, int]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1569
        max_allowed = 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1570
        if current_line.depth == 0:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1571
            max_allowed = 1 if self.is_pyi else 2
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1572
        if current_line.leaves:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1573
            # Consume the first leaf's extra newlines.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1574
            first_leaf = current_line.leaves[0]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1575
            before = first_leaf.prefix.count("\n")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1576
            before = min(before, max_allowed)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1577
            first_leaf.prefix = ""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1578
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1579
            before = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1580
        depth = current_line.depth
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1581
        while self.previous_defs and self.previous_defs[-1] >= depth:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1582
            self.previous_defs.pop()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1583
            if self.is_pyi:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1584
                before = 0 if depth else 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1585
            else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1586
                before = 1 if depth else 2
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1587
        if current_line.is_decorator or current_line.is_def or current_line.is_class:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1588
            return self._maybe_empty_lines_for_class_or_def(current_line, before)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1589
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1590
        if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1591
            self.previous_line
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1592
            and self.previous_line.is_import
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1593
            and not current_line.is_import
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1594
            and depth == self.previous_line.depth
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1595
        ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1596
            return (before or 1), 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1597
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1598
        if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1599
            self.previous_line
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1600
            and self.previous_line.is_class
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1601
            and current_line.is_triple_quoted_string
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1602
        ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1603
            return before, 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1604
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1605
        return before, 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1606
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1607
    def _maybe_empty_lines_for_class_or_def(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1608
        self, current_line: Line, before: int
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1609
    ) -> Tuple[int, int]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1610
        if not current_line.is_decorator:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1611
            self.previous_defs.append(current_line.depth)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1612
        if self.previous_line is None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1613
            # Don't insert empty lines before the first line in the file.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1614
            return 0, 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1615
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1616
        if self.previous_line.is_decorator:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1617
            return 0, 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1618
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1619
        if self.previous_line.depth < current_line.depth and (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1620
            self.previous_line.is_class or self.previous_line.is_def
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1621
        ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1622
            return 0, 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1623
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1624
        if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1625
            self.previous_line.is_comment
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1626
            and self.previous_line.depth == current_line.depth
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1627
            and before == 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1628
        ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1629
            return 0, 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1630
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1631
        if self.is_pyi:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1632
            if self.previous_line.depth > current_line.depth:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1633
                newlines = 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1634
            elif current_line.is_class or self.previous_line.is_class:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1635
                if current_line.is_stub_class and self.previous_line.is_stub_class:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1636
                    # No blank line between classes with an empty body
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1637
                    newlines = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1638
                else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1639
                    newlines = 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1640
            elif current_line.is_def and not self.previous_line.is_def:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1641
                # Blank line between a block of functions and a block of non-functions
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1642
                newlines = 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1643
            else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1644
                newlines = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1645
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1646
            newlines = 2
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1647
        if current_line.depth and newlines:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1648
            newlines -= 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1649
        return newlines, 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1650
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1651
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1652
@dataclass
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1653
class LineGenerator(Visitor[Line]):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1654
    """Generates reformatted Line objects.  Empty lines are not emitted.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1655
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1656
    Note: destroys the tree it's visiting by mutating prefixes of its leaves
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1657
    in ways that will no longer stringify to valid Python code on the tree.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1658
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1659
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1660
    is_pyi: bool = False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1661
    normalize_strings: bool = True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1662
    current_line: Line = Factory(Line)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1663
    remove_u_prefix: bool = False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1664
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1665
    def line(self, indent: int = 0) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1666
        """Generate a line.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1667
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1668
        If the line is empty, only emit if it makes sense.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1669
        If the line is too long, split it first and then generate.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1670
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1671
        If any lines were generated, set up a new current_line.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1672
        """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1673
        if not self.current_line:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1674
            self.current_line.depth += indent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1675
            return  # Line is empty, don't emit. Creating a new one unnecessary.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1676
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1677
        complete_line = self.current_line
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1678
        self.current_line = Line(depth=complete_line.depth + indent)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1679
        yield complete_line
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1680
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1681
    def visit_default(self, node: LN) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1682
        """Default `visit_*()` implementation. Recurses to children of `node`."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1683
        if isinstance(node, Leaf):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1684
            any_open_brackets = self.current_line.bracket_tracker.any_open_brackets()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1685
            for comment in generate_comments(node):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1686
                if any_open_brackets:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1687
                    # any comment within brackets is subject to splitting
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1688
                    self.current_line.append(comment)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1689
                elif comment.type == token.COMMENT:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1690
                    # regular trailing comment
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1691
                    self.current_line.append(comment)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1692
                    yield from self.line()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1693
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1694
                else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1695
                    # regular standalone comment
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1696
                    yield from self.line()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1697
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1698
                    self.current_line.append(comment)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1699
                    yield from self.line()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1700
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1701
            normalize_prefix(node, inside_brackets=any_open_brackets)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1702
            if self.normalize_strings and node.type == token.STRING:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1703
                normalize_string_prefix(node, remove_u_prefix=self.remove_u_prefix)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1704
                normalize_string_quotes(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1705
            if node.type == token.NUMBER:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1706
                normalize_numeric_literal(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1707
            if node.type not in WHITESPACE:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1708
                self.current_line.append(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1709
        yield from super().visit_default(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1710
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1711
    def visit_atom(self, node: Node) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1712
        # Always make parentheses invisible around a single node, because it should
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1713
        # not be needed (except in the case of yield, where removing the parentheses
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1714
        # produces a SyntaxError).
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1715
        if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1716
            len(node.children) == 3
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1717
            and isinstance(node.children[0], Leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1718
            and node.children[0].type == token.LPAR
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1719
            and isinstance(node.children[2], Leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1720
            and node.children[2].type == token.RPAR
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1721
            and isinstance(node.children[1], Leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1722
            and not (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1723
                node.children[1].type == token.NAME
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1724
                and node.children[1].value == "yield"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1725
            )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1726
        ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1727
            node.children[0].value = ""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1728
            node.children[2].value = ""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1729
        yield from super().visit_default(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1730
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1731
    def visit_factor(self, node: Node) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1732
        """Force parentheses between a unary op and a binary power:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1733
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1734
        -2 ** 8 -> -(2 ** 8)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1735
        """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1736
        child = node.children[1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1737
        if child.type == syms.power and len(child.children) == 3:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1738
            lpar = Leaf(token.LPAR, "(")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1739
            rpar = Leaf(token.RPAR, ")")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1740
            index = child.remove() or 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1741
            node.insert_child(index, Node(syms.atom, [lpar, child, rpar]))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1742
        yield from self.visit_default(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1743
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1744
    def visit_INDENT(self, node: Node) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1745
        """Increase indentation level, maybe yield a line."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1746
        # In blib2to3 INDENT never holds comments.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1747
        yield from self.line(+1)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1748
        yield from self.visit_default(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1749
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1750
    def visit_DEDENT(self, node: Node) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1751
        """Decrease indentation level, maybe yield a line."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1752
        # The current line might still wait for trailing comments.  At DEDENT time
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1753
        # there won't be any (they would be prefixes on the preceding NEWLINE).
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1754
        # Emit the line then.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1755
        yield from self.line()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1756
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1757
        # While DEDENT has no value, its prefix may contain standalone comments
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1758
        # that belong to the current indentation level.  Get 'em.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1759
        yield from self.visit_default(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1760
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1761
        # Finally, emit the dedent.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1762
        yield from self.line(-1)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1763
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1764
    def visit_stmt(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1765
        self, node: Node, keywords: Set[str], parens: Set[str]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1766
    ) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1767
        """Visit a statement.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1768
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1769
        This implementation is shared for `if`, `while`, `for`, `try`, `except`,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1770
        `def`, `with`, `class`, `assert` and assignments.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1771
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1772
        The relevant Python language `keywords` for a given statement will be
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1773
        NAME leaves within it. This methods puts those on a separate line.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1774
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1775
        `parens` holds a set of string leaf values immediately after which
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1776
        invisible parens should be put.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1777
        """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1778
        normalize_invisible_parens(node, parens_after=parens)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1779
        for child in node.children:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1780
            if child.type == token.NAME and child.value in keywords:  # type: ignore
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1781
                yield from self.line()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1782
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1783
            yield from self.visit(child)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1784
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1785
    def visit_suite(self, node: Node) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1786
        """Visit a suite."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1787
        if self.is_pyi and is_stub_suite(node):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1788
            yield from self.visit(node.children[2])
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1789
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1790
            yield from self.visit_default(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1791
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1792
    def visit_simple_stmt(self, node: Node) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1793
        """Visit a statement without nested statements."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1794
        is_suite_like = node.parent and node.parent.type in STATEMENT
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1795
        if is_suite_like:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1796
            if self.is_pyi and is_stub_body(node):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1797
                yield from self.visit_default(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1798
            else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1799
                yield from self.line(+1)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1800
                yield from self.visit_default(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1801
                yield from self.line(-1)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1802
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1803
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1804
            if not self.is_pyi or not node.parent or not is_stub_suite(node.parent):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1805
                yield from self.line()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1806
            yield from self.visit_default(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1807
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1808
    def visit_async_stmt(self, node: Node) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1809
        """Visit `async def`, `async for`, `async with`."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1810
        yield from self.line()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1811
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1812
        children = iter(node.children)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1813
        for child in children:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1814
            yield from self.visit(child)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1815
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1816
            if child.type == token.ASYNC:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1817
                break
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1818
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1819
        internal_stmt = next(children)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1820
        for child in internal_stmt.children:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1821
            yield from self.visit(child)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1822
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1823
    def visit_decorators(self, node: Node) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1824
        """Visit decorators."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1825
        for child in node.children:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1826
            yield from self.line()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1827
            yield from self.visit(child)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1828
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1829
    def visit_SEMI(self, leaf: Leaf) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1830
        """Remove a semicolon and put the other statement on a separate line."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1831
        yield from self.line()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1832
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1833
    def visit_ENDMARKER(self, leaf: Leaf) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1834
        """End of file. Process outstanding comments and end with a newline."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1835
        yield from self.visit_default(leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1836
        yield from self.line()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1837
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1838
    def visit_STANDALONE_COMMENT(self, leaf: Leaf) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1839
        if not self.current_line.bracket_tracker.any_open_brackets():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1840
            yield from self.line()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1841
        yield from self.visit_default(leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1842
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1843
    def __attrs_post_init__(self) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1844
        """You are in a twisty little maze of passages."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1845
        v = self.visit_stmt
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1846
        Ø: Set[str] = set()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1847
        self.visit_assert_stmt = partial(v, keywords={"assert"}, parens={"assert", ","})
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1848
        self.visit_if_stmt = partial(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1849
            v, keywords={"if", "else", "elif"}, parens={"if", "elif"}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1850
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1851
        self.visit_while_stmt = partial(v, keywords={"while", "else"}, parens={"while"})
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1852
        self.visit_for_stmt = partial(v, keywords={"for", "else"}, parens={"for", "in"})
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1853
        self.visit_try_stmt = partial(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1854
            v, keywords={"try", "except", "else", "finally"}, parens=Ø
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1855
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1856
        self.visit_except_clause = partial(v, keywords={"except"}, parens=Ø)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1857
        self.visit_with_stmt = partial(v, keywords={"with"}, parens=Ø)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1858
        self.visit_funcdef = partial(v, keywords={"def"}, parens=Ø)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1859
        self.visit_classdef = partial(v, keywords={"class"}, parens=Ø)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1860
        self.visit_expr_stmt = partial(v, keywords=Ø, parens=ASSIGNMENTS)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1861
        self.visit_return_stmt = partial(v, keywords={"return"}, parens={"return"})
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1862
        self.visit_import_from = partial(v, keywords=Ø, parens={"import"})
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1863
        self.visit_del_stmt = partial(v, keywords=Ø, parens={"del"})
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1864
        self.visit_async_funcdef = self.visit_async_stmt
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1865
        self.visit_decorated = self.visit_decorators
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1866
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1867
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1868
IMPLICIT_TUPLE = {syms.testlist, syms.testlist_star_expr, syms.exprlist}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1869
BRACKET = {token.LPAR: token.RPAR, token.LSQB: token.RSQB, token.LBRACE: token.RBRACE}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1870
OPENING_BRACKETS = set(BRACKET.keys())
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1871
CLOSING_BRACKETS = set(BRACKET.values())
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1872
BRACKETS = OPENING_BRACKETS | CLOSING_BRACKETS
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1873
ALWAYS_NO_SPACE = CLOSING_BRACKETS | {token.COMMA, STANDALONE_COMMENT}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1874
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1875
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1876
def whitespace(leaf: Leaf, *, complex_subscript: bool) -> str:  # noqa: C901
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1877
    """Return whitespace prefix if needed for the given `leaf`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1878
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1879
    `complex_subscript` signals whether the given leaf is part of a subscription
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1880
    which has non-trivial arguments, like arithmetic expressions or function calls.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1881
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1882
    NO = ""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1883
    SPACE = " "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1884
    DOUBLESPACE = "  "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1885
    t = leaf.type
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1886
    p = leaf.parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1887
    v = leaf.value
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1888
    if t in ALWAYS_NO_SPACE:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1889
        return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1890
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1891
    if t == token.COMMENT:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1892
        return DOUBLESPACE
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1893
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1894
    assert p is not None, f"INTERNAL ERROR: hand-made leaf without parent: {leaf!r}"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1895
    if t == token.COLON and p.type not in {
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1896
        syms.subscript,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1897
        syms.subscriptlist,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1898
        syms.sliceop,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1899
    }:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1900
        return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1901
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1902
    prev = leaf.prev_sibling
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1903
    if not prev:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1904
        prevp = preceding_leaf(p)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1905
        if not prevp or prevp.type in OPENING_BRACKETS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1906
            return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1907
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1908
        if t == token.COLON:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1909
            if prevp.type == token.COLON:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1910
                return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1911
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1912
            elif prevp.type != token.COMMA and not complex_subscript:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1913
                return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1914
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1915
            return SPACE
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1916
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1917
        if prevp.type == token.EQUAL:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1918
            if prevp.parent:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1919
                if prevp.parent.type in {
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1920
                    syms.arglist,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1921
                    syms.argument,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1922
                    syms.parameters,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1923
                    syms.varargslist,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1924
                }:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1925
                    return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1926
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1927
                elif prevp.parent.type == syms.typedargslist:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1928
                    # A bit hacky: if the equal sign has whitespace, it means we
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1929
                    # previously found it's a typed argument.  So, we're using
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1930
                    # that, too.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1931
                    return prevp.prefix
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1932
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1933
        elif prevp.type in VARARGS_SPECIALS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1934
            if is_vararg(prevp, within=VARARGS_PARENTS | UNPACKING_PARENTS):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1935
                return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1936
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1937
        elif prevp.type == token.COLON:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1938
            if prevp.parent and prevp.parent.type in {syms.subscript, syms.sliceop}:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1939
                return SPACE if complex_subscript else NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1940
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1941
        elif (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1942
            prevp.parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1943
            and prevp.parent.type == syms.factor
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1944
            and prevp.type in MATH_OPERATORS
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1945
        ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1946
            return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1947
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1948
        elif (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1949
            prevp.type == token.RIGHTSHIFT
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1950
            and prevp.parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1951
            and prevp.parent.type == syms.shift_expr
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1952
            and prevp.prev_sibling
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1953
            and prevp.prev_sibling.type == token.NAME
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1954
            and prevp.prev_sibling.value == "print"  # type: ignore
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1955
        ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1956
            # Python 2 print chevron
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1957
            return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1958
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1959
    elif prev.type in OPENING_BRACKETS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1960
        return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1961
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1962
    if p.type in {syms.parameters, syms.arglist}:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1963
        # untyped function signatures or calls
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1964
        if not prev or prev.type != token.COMMA:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1965
            return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1966
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1967
    elif p.type == syms.varargslist:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1968
        # lambdas
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1969
        if prev and prev.type != token.COMMA:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1970
            return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1971
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1972
    elif p.type == syms.typedargslist:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1973
        # typed function signatures
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1974
        if not prev:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1975
            return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1976
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1977
        if t == token.EQUAL:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1978
            if prev.type != syms.tname:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1979
                return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1980
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1981
        elif prev.type == token.EQUAL:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1982
            # A bit hacky: if the equal sign has whitespace, it means we
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1983
            # previously found it's a typed argument.  So, we're using that, too.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1984
            return prev.prefix
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1985
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1986
        elif prev.type != token.COMMA:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1987
            return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1988
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1989
    elif p.type == syms.tname:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1990
        # type names
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1991
        if not prev:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1992
            prevp = preceding_leaf(p)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1993
            if not prevp or prevp.type != token.COMMA:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1994
                return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1995
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1996
    elif p.type == syms.trailer:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1997
        # attributes and calls
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1998
        if t == token.LPAR or t == token.RPAR:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  1999
            return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2000
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2001
        if not prev:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2002
            if t == token.DOT:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2003
                prevp = preceding_leaf(p)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2004
                if not prevp or prevp.type != token.NUMBER:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2005
                    return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2006
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2007
            elif t == token.LSQB:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2008
                return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2009
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2010
        elif prev.type != token.COMMA:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2011
            return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2012
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2013
    elif p.type == syms.argument:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2014
        # single argument
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2015
        if t == token.EQUAL:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2016
            return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2017
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2018
        if not prev:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2019
            prevp = preceding_leaf(p)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2020
            if not prevp or prevp.type == token.LPAR:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2021
                return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2022
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2023
        elif prev.type in {token.EQUAL} | VARARGS_SPECIALS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2024
            return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2025
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2026
    elif p.type == syms.decorator:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2027
        # decorators
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2028
        return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2029
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2030
    elif p.type == syms.dotted_name:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2031
        if prev:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2032
            return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2033
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2034
        prevp = preceding_leaf(p)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2035
        if not prevp or prevp.type == token.AT or prevp.type == token.DOT:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2036
            return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2037
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2038
    elif p.type == syms.classdef:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2039
        if t == token.LPAR:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2040
            return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2041
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2042
        if prev and prev.type == token.LPAR:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2043
            return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2044
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2045
    elif p.type in {syms.subscript, syms.sliceop}:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2046
        # indexing
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2047
        if not prev:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2048
            assert p.parent is not None, "subscripts are always parented"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2049
            if p.parent.type == syms.subscriptlist:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2050
                return SPACE
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2051
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2052
            return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2053
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2054
        elif not complex_subscript:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2055
            return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2056
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2057
    elif p.type == syms.atom:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2058
        if prev and t == token.DOT:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2059
            # dots, but not the first one.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2060
            return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2061
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2062
    elif p.type == syms.dictsetmaker:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2063
        # dict unpacking
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2064
        if prev and prev.type == token.DOUBLESTAR:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2065
            return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2066
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2067
    elif p.type in {syms.factor, syms.star_expr}:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2068
        # unary ops
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2069
        if not prev:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2070
            prevp = preceding_leaf(p)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2071
            if not prevp or prevp.type in OPENING_BRACKETS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2072
                return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2073
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2074
            prevp_parent = prevp.parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2075
            assert prevp_parent is not None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2076
            if prevp.type == token.COLON and prevp_parent.type in {
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2077
                syms.subscript,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2078
                syms.sliceop,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2079
            }:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2080
                return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2081
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2082
            elif prevp.type == token.EQUAL and prevp_parent.type == syms.argument:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2083
                return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2084
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2085
        elif t in {token.NAME, token.NUMBER, token.STRING}:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2086
            return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2087
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2088
    elif p.type == syms.import_from:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2089
        if t == token.DOT:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2090
            if prev and prev.type == token.DOT:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2091
                return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2092
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2093
        elif t == token.NAME:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2094
            if v == "import":
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2095
                return SPACE
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2096
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2097
            if prev and prev.type == token.DOT:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2098
                return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2099
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2100
    elif p.type == syms.sliceop:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2101
        return NO
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2102
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2103
    return SPACE
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2104
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2105
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2106
def preceding_leaf(node: Optional[LN]) -> Optional[Leaf]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2107
    """Return the first leaf that precedes `node`, if any."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2108
    while node:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2109
        res = node.prev_sibling
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2110
        if res:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2111
            if isinstance(res, Leaf):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2112
                return res
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2113
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2114
            try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2115
                return list(res.leaves())[-1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2116
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2117
            except IndexError:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2118
                return None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2119
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2120
        node = node.parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2121
    return None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2122
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2123
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2124
def child_towards(ancestor: Node, descendant: LN) -> Optional[LN]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2125
    """Return the child of `ancestor` that contains `descendant`."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2126
    node: Optional[LN] = descendant
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2127
    while node and node.parent != ancestor:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2128
        node = node.parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2129
    return node
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2130
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2131
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2132
def container_of(leaf: Leaf) -> LN:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2133
    """Return `leaf` or one of its ancestors that is the topmost container of it.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2134
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2135
    By "container" we mean a node where `leaf` is the very first child.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2136
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2137
    same_prefix = leaf.prefix
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2138
    container: LN = leaf
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2139
    while container:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2140
        parent = container.parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2141
        if parent is None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2142
            break
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2143
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2144
        if parent.children[0].prefix != same_prefix:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2145
            break
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2146
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2147
        if parent.type == syms.file_input:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2148
            break
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2149
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2150
        if parent.prev_sibling is not None and parent.prev_sibling.type in BRACKETS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2151
            break
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2152
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2153
        container = parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2154
    return container
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2155
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2156
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2157
def is_split_after_delimiter(leaf: Leaf, previous: Optional[Leaf] = None) -> Priority:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2158
    """Return the priority of the `leaf` delimiter, given a line break after it.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2159
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2160
    The delimiter priorities returned here are from those delimiters that would
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2161
    cause a line break after themselves.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2162
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2163
    Higher numbers are higher priority.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2164
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2165
    if leaf.type == token.COMMA:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2166
        return COMMA_PRIORITY
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2167
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2168
    return 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2169
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2170
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2171
def is_split_before_delimiter(leaf: Leaf, previous: Optional[Leaf] = None) -> Priority:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2172
    """Return the priority of the `leaf` delimiter, given a line break before it.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2173
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2174
    The delimiter priorities returned here are from those delimiters that would
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2175
    cause a line break before themselves.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2176
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2177
    Higher numbers are higher priority.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2178
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2179
    if is_vararg(leaf, within=VARARGS_PARENTS | UNPACKING_PARENTS):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2180
        # * and ** might also be MATH_OPERATORS but in this case they are not.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2181
        # Don't treat them as a delimiter.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2182
        return 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2183
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2184
    if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2185
        leaf.type == token.DOT
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2186
        and leaf.parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2187
        and leaf.parent.type not in {syms.import_from, syms.dotted_name}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2188
        and (previous is None or previous.type in CLOSING_BRACKETS)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2189
    ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2190
        return DOT_PRIORITY
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2191
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2192
    if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2193
        leaf.type in MATH_OPERATORS
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2194
        and leaf.parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2195
        and leaf.parent.type not in {syms.factor, syms.star_expr}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2196
    ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2197
        return MATH_PRIORITIES[leaf.type]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2198
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2199
    if leaf.type in COMPARATORS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2200
        return COMPARATOR_PRIORITY
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2201
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2202
    if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2203
        leaf.type == token.STRING
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2204
        and previous is not None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2205
        and previous.type == token.STRING
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2206
    ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2207
        return STRING_PRIORITY
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2208
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2209
    if leaf.type not in {token.NAME, token.ASYNC}:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2210
        return 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2211
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2212
    if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2213
        leaf.value == "for"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2214
        and leaf.parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2215
        and leaf.parent.type in {syms.comp_for, syms.old_comp_for}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2216
        or leaf.type == token.ASYNC
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2217
    ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2218
        if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2219
            not isinstance(leaf.prev_sibling, Leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2220
            or leaf.prev_sibling.value != "async"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2221
        ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2222
            return COMPREHENSION_PRIORITY
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2223
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2224
    if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2225
        leaf.value == "if"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2226
        and leaf.parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2227
        and leaf.parent.type in {syms.comp_if, syms.old_comp_if}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2228
    ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2229
        return COMPREHENSION_PRIORITY
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2230
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2231
    if leaf.value in {"if", "else"} and leaf.parent and leaf.parent.type == syms.test:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2232
        return TERNARY_PRIORITY
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2233
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2234
    if leaf.value == "is":
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2235
        return COMPARATOR_PRIORITY
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2236
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2237
    if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2238
        leaf.value == "in"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2239
        and leaf.parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2240
        and leaf.parent.type in {syms.comp_op, syms.comparison}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2241
        and not (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2242
            previous is not None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2243
            and previous.type == token.NAME
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2244
            and previous.value == "not"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2245
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2246
    ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2247
        return COMPARATOR_PRIORITY
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2248
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2249
    if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2250
        leaf.value == "not"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2251
        and leaf.parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2252
        and leaf.parent.type == syms.comp_op
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2253
        and not (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2254
            previous is not None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2255
            and previous.type == token.NAME
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2256
            and previous.value == "is"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2257
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2258
    ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2259
        return COMPARATOR_PRIORITY
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2260
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2261
    if leaf.value in LOGIC_OPERATORS and leaf.parent:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2262
        return LOGIC_PRIORITY
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2263
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2264
    return 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2265
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2266
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2267
FMT_OFF = {"# fmt: off", "# fmt:off", "# yapf: disable"}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2268
FMT_ON = {"# fmt: on", "# fmt:on", "# yapf: enable"}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2269
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2270
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2271
def generate_comments(leaf: LN) -> Iterator[Leaf]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2272
    """Clean the prefix of the `leaf` and generate comments from it, if any.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2273
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2274
    Comments in lib2to3 are shoved into the whitespace prefix.  This happens
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2275
    in `pgen2/driver.py:Driver.parse_tokens()`.  This was a brilliant implementation
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2276
    move because it does away with modifying the grammar to include all the
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2277
    possible places in which comments can be placed.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2278
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2279
    The sad consequence for us though is that comments don't "belong" anywhere.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2280
    This is why this function generates simple parentless Leaf objects for
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2281
    comments.  We simply don't know what the correct parent should be.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2282
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2283
    No matter though, we can live without this.  We really only need to
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2284
    differentiate between inline and standalone comments.  The latter don't
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2285
    share the line with any code.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2286
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2287
    Inline comments are emitted as regular token.COMMENT leaves.  Standalone
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2288
    are emitted with a fake STANDALONE_COMMENT token identifier.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2289
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2290
    for pc in list_comments(leaf.prefix, is_endmarker=leaf.type == token.ENDMARKER):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2291
        yield Leaf(pc.type, pc.value, prefix="\n" * pc.newlines)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2292
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2293
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2294
@dataclass
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2295
class ProtoComment:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2296
    """Describes a piece of syntax that is a comment.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2297
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2298
    It's not a :class:`blib2to3.pytree.Leaf` so that:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2299
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2300
    * it can be cached (`Leaf` objects should not be reused more than once as
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2301
      they store their lineno, column, prefix, and parent information);
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2302
    * `newlines` and `consumed` fields are kept separate from the `value`. This
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2303
      simplifies handling of special marker comments like ``# fmt: off/on``.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2304
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2305
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2306
    type: int  # token.COMMENT or STANDALONE_COMMENT
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2307
    value: str  # content of the comment
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2308
    newlines: int  # how many newlines before the comment
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2309
    consumed: int  # how many characters of the original leaf's prefix did we consume
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2310
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2311
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2312
@lru_cache(maxsize=4096)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2313
def list_comments(prefix: str, *, is_endmarker: bool) -> List[ProtoComment]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2314
    """Return a list of :class:`ProtoComment` objects parsed from the given `prefix`."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2315
    result: List[ProtoComment] = []
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2316
    if not prefix or "#" not in prefix:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2317
        return result
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2318
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2319
    consumed = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2320
    nlines = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2321
    ignored_lines = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2322
    for index, line in enumerate(prefix.split("\n")):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2323
        consumed += len(line) + 1  # adding the length of the split '\n'
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2324
        line = line.lstrip()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2325
        if not line:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2326
            nlines += 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2327
        if not line.startswith("#"):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2328
            # Escaped newlines outside of a comment are not really newlines at
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2329
            # all. We treat a single-line comment following an escaped newline
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2330
            # as a simple trailing comment.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2331
            if line.endswith("\\"):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2332
                ignored_lines += 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2333
            continue
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2334
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2335
        if index == ignored_lines and not is_endmarker:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2336
            comment_type = token.COMMENT  # simple trailing comment
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2337
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2338
            comment_type = STANDALONE_COMMENT
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2339
        comment = make_comment(line)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2340
        result.append(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2341
            ProtoComment(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2342
                type=comment_type, value=comment, newlines=nlines, consumed=consumed
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2343
            )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2344
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2345
        nlines = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2346
    return result
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2347
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2348
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2349
def make_comment(content: str) -> str:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2350
    """Return a consistently formatted comment from the given `content` string.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2351
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2352
    All comments (except for "##", "#!", "#:", '#'", "#%%") should have a single
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2353
    space between the hash sign and the content.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2354
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2355
    If `content` didn't start with a hash sign, one is provided.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2356
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2357
    content = content.rstrip()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2358
    if not content:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2359
        return "#"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2360
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2361
    if content[0] == "#":
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2362
        content = content[1:]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2363
    if content and content[0] not in " !:#'%":
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2364
        content = " " + content
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2365
    return "#" + content
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2366
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2367
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2368
def split_line(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2369
    line: Line,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2370
    line_length: int,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2371
    inner: bool = False,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2372
    features: Collection[Feature] = (),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2373
) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2374
    """Split a `line` into potentially many lines.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2375
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2376
    They should fit in the allotted `line_length` but might not be able to.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2377
    `inner` signifies that there were a pair of brackets somewhere around the
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2378
    current `line`, possibly transitively. This means we can fallback to splitting
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2379
    by delimiters if the LHS/RHS don't yield any results.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2380
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2381
    `features` are syntactical features that may be used in the output.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2382
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2383
    if line.is_comment:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2384
        yield line
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2385
        return
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2386
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2387
    line_str = str(line).strip("\n")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2388
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2389
    if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2390
        not line.contains_uncollapsable_type_comments()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2391
        and not line.should_explode
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2392
        and not line.is_collection_with_optional_trailing_comma
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2393
        and (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2394
            is_line_short_enough(line, line_length=line_length, line_str=line_str)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2395
            or line.contains_unsplittable_type_ignore()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2396
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2397
    ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2398
        yield line
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2399
        return
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2400
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2401
    split_funcs: List[SplitFunc]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2402
    if line.is_def:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2403
        split_funcs = [left_hand_split]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2404
    else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2405
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2406
        def rhs(line: Line, features: Collection[Feature]) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2407
            for omit in generate_trailers_to_omit(line, line_length):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2408
                lines = list(right_hand_split(line, line_length, features, omit=omit))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2409
                if is_line_short_enough(lines[0], line_length=line_length):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2410
                    yield from lines
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2411
                    return
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2412
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2413
            # All splits failed, best effort split with no omits.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2414
            # This mostly happens to multiline strings that are by definition
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2415
            # reported as not fitting a single line.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2416
            yield from right_hand_split(line, line_length, features=features)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2417
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2418
        if line.inside_brackets:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2419
            split_funcs = [delimiter_split, standalone_comment_split, rhs]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2420
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2421
            split_funcs = [rhs]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2422
    for split_func in split_funcs:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2423
        # We are accumulating lines in `result` because we might want to abort
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2424
        # mission and return the original line in the end, or attempt a different
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2425
        # split altogether.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2426
        result: List[Line] = []
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2427
        try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2428
            for l in split_func(line, features):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2429
                if str(l).strip("\n") == line_str:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2430
                    raise CannotSplit("Split function returned an unchanged result")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2431
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2432
                result.extend(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2433
                    split_line(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2434
                        l, line_length=line_length, inner=True, features=features
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2435
                    )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2436
                )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2437
        except CannotSplit:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2438
            continue
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2439
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2440
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2441
            yield from result
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2442
            break
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2443
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2444
    else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2445
        yield line
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2446
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2447
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2448
def left_hand_split(line: Line, features: Collection[Feature] = ()) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2449
    """Split line into many lines, starting with the first matching bracket pair.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2450
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2451
    Note: this usually looks weird, only use this for function definitions.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2452
    Prefer RHS otherwise.  This is why this function is not symmetrical with
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2453
    :func:`right_hand_split` which also handles optional parentheses.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2454
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2455
    tail_leaves: List[Leaf] = []
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2456
    body_leaves: List[Leaf] = []
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2457
    head_leaves: List[Leaf] = []
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2458
    current_leaves = head_leaves
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2459
    matching_bracket = None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2460
    for leaf in line.leaves:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2461
        if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2462
            current_leaves is body_leaves
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2463
            and leaf.type in CLOSING_BRACKETS
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2464
            and leaf.opening_bracket is matching_bracket
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2465
        ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2466
            current_leaves = tail_leaves if body_leaves else head_leaves
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2467
        current_leaves.append(leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2468
        if current_leaves is head_leaves:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2469
            if leaf.type in OPENING_BRACKETS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2470
                matching_bracket = leaf
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2471
                current_leaves = body_leaves
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2472
    if not matching_bracket:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2473
        raise CannotSplit("No brackets found")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2474
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2475
    head = bracket_split_build_line(head_leaves, line, matching_bracket)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2476
    body = bracket_split_build_line(body_leaves, line, matching_bracket, is_body=True)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2477
    tail = bracket_split_build_line(tail_leaves, line, matching_bracket)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2478
    bracket_split_succeeded_or_raise(head, body, tail)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2479
    for result in (head, body, tail):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2480
        if result:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2481
            yield result
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2482
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2483
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2484
def right_hand_split(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2485
    line: Line,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2486
    line_length: int,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2487
    features: Collection[Feature] = (),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2488
    omit: Collection[LeafID] = (),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2489
) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2490
    """Split line into many lines, starting with the last matching bracket pair.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2491
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2492
    If the split was by optional parentheses, attempt splitting without them, too.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2493
    `omit` is a collection of closing bracket IDs that shouldn't be considered for
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2494
    this split.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2495
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2496
    Note: running this function modifies `bracket_depth` on the leaves of `line`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2497
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2498
    tail_leaves: List[Leaf] = []
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2499
    body_leaves: List[Leaf] = []
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2500
    head_leaves: List[Leaf] = []
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2501
    current_leaves = tail_leaves
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2502
    opening_bracket = None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2503
    closing_bracket = None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2504
    for leaf in reversed(line.leaves):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2505
        if current_leaves is body_leaves:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2506
            if leaf is opening_bracket:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2507
                current_leaves = head_leaves if body_leaves else tail_leaves
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2508
        current_leaves.append(leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2509
        if current_leaves is tail_leaves:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2510
            if leaf.type in CLOSING_BRACKETS and id(leaf) not in omit:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2511
                opening_bracket = leaf.opening_bracket
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2512
                closing_bracket = leaf
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2513
                current_leaves = body_leaves
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2514
    if not (opening_bracket and closing_bracket and head_leaves):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2515
        # If there is no opening or closing_bracket that means the split failed and
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2516
        # all content is in the tail.  Otherwise, if `head_leaves` are empty, it means
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2517
        # the matching `opening_bracket` wasn't available on `line` anymore.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2518
        raise CannotSplit("No brackets found")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2519
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2520
    tail_leaves.reverse()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2521
    body_leaves.reverse()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2522
    head_leaves.reverse()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2523
    head = bracket_split_build_line(head_leaves, line, opening_bracket)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2524
    body = bracket_split_build_line(body_leaves, line, opening_bracket, is_body=True)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2525
    tail = bracket_split_build_line(tail_leaves, line, opening_bracket)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2526
    bracket_split_succeeded_or_raise(head, body, tail)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2527
    if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2528
        # the body shouldn't be exploded
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2529
        not body.should_explode
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2530
        # the opening bracket is an optional paren
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2531
        and opening_bracket.type == token.LPAR
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2532
        and not opening_bracket.value
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2533
        # the closing bracket is an optional paren
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2534
        and closing_bracket.type == token.RPAR
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2535
        and not closing_bracket.value
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2536
        # it's not an import (optional parens are the only thing we can split on
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2537
        # in this case; attempting a split without them is a waste of time)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2538
        and not line.is_import
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2539
        # there are no standalone comments in the body
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2540
        and not body.contains_standalone_comments(0)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2541
        # and we can actually remove the parens
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2542
        and can_omit_invisible_parens(body, line_length)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2543
    ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2544
        omit = {id(closing_bracket), *omit}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2545
        try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2546
            yield from right_hand_split(line, line_length, features=features, omit=omit)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2547
            return
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2548
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2549
        except CannotSplit:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2550
            if not (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2551
                can_be_split(body)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2552
                or is_line_short_enough(body, line_length=line_length)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2553
            ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2554
                raise CannotSplit(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2555
                    "Splitting failed, body is still too long and can't be split."
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2556
                )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2557
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2558
            elif head.contains_multiline_strings() or tail.contains_multiline_strings():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2559
                raise CannotSplit(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2560
                    "The current optional pair of parentheses is bound to fail to "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2561
                    "satisfy the splitting algorithm because the head or the tail "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2562
                    "contains multiline strings which by definition never fit one "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2563
                    "line."
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2564
                )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2565
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2566
    ensure_visible(opening_bracket)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2567
    ensure_visible(closing_bracket)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2568
    for result in (head, body, tail):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2569
        if result:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2570
            yield result
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2571
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2572
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2573
def bracket_split_succeeded_or_raise(head: Line, body: Line, tail: Line) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2574
    """Raise :exc:`CannotSplit` if the last left- or right-hand split failed.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2575
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2576
    Do nothing otherwise.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2577
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2578
    A left- or right-hand split is based on a pair of brackets. Content before
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2579
    (and including) the opening bracket is left on one line, content inside the
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2580
    brackets is put on a separate line, and finally content starting with and
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2581
    following the closing bracket is put on a separate line.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2582
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2583
    Those are called `head`, `body`, and `tail`, respectively. If the split
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2584
    produced the same line (all content in `head`) or ended up with an empty `body`
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2585
    and the `tail` is just the closing bracket, then it's considered failed.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2586
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2587
    tail_len = len(str(tail).strip())
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2588
    if not body:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2589
        if tail_len == 0:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2590
            raise CannotSplit("Splitting brackets produced the same line")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2591
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2592
        elif tail_len < 3:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2593
            raise CannotSplit(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2594
                f"Splitting brackets on an empty body to save "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2595
                f"{tail_len} characters is not worth it"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2596
            )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2597
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2598
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2599
def bracket_split_build_line(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2600
    leaves: List[Leaf], original: Line, opening_bracket: Leaf, *, is_body: bool = False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2601
) -> Line:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2602
    """Return a new line with given `leaves` and respective comments from `original`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2603
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2604
    If `is_body` is True, the result line is one-indented inside brackets and as such
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2605
    has its first leaf's prefix normalized and a trailing comma added when expected.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2606
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2607
    result = Line(depth=original.depth)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2608
    if is_body:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2609
        result.inside_brackets = True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2610
        result.depth += 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2611
        if leaves:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2612
            # Since body is a new indent level, remove spurious leading whitespace.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2613
            normalize_prefix(leaves[0], inside_brackets=True)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2614
            # Ensure a trailing comma for imports and standalone function arguments, but
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2615
            # be careful not to add one after any comments.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2616
            no_commas = original.is_def and not any(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2617
                l.type == token.COMMA for l in leaves
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2618
            )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2619
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2620
            if original.is_import or no_commas:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2621
                for i in range(len(leaves) - 1, -1, -1):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2622
                    if leaves[i].type == STANDALONE_COMMENT:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2623
                        continue
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2624
                    elif leaves[i].type == token.COMMA:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2625
                        break
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2626
                    else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2627
                        leaves.insert(i + 1, Leaf(token.COMMA, ","))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2628
                        break
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2629
    # Populate the line
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2630
    for leaf in leaves:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2631
        result.append(leaf, preformatted=True)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2632
        for comment_after in original.comments_after(leaf):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2633
            result.append(comment_after, preformatted=True)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2634
    if is_body:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2635
        result.should_explode = should_explode(result, opening_bracket)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2636
    return result
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2637
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2638
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2639
def dont_increase_indentation(split_func: SplitFunc) -> SplitFunc:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2640
    """Normalize prefix of the first leaf in every line returned by `split_func`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2641
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2642
    This is a decorator over relevant split functions.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2643
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2644
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2645
    @wraps(split_func)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2646
    def split_wrapper(line: Line, features: Collection[Feature] = ()) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2647
        for l in split_func(line, features):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2648
            normalize_prefix(l.leaves[0], inside_brackets=True)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2649
            yield l
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2650
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2651
    return split_wrapper
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2652
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2653
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2654
@dont_increase_indentation
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2655
def delimiter_split(line: Line, features: Collection[Feature] = ()) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2656
    """Split according to delimiters of the highest priority.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2657
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2658
    If the appropriate Features are given, the split will add trailing commas
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2659
    also in function signatures and calls that contain `*` and `**`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2660
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2661
    try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2662
        last_leaf = line.leaves[-1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2663
    except IndexError:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2664
        raise CannotSplit("Line empty")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2665
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2666
    bt = line.bracket_tracker
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2667
    try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2668
        delimiter_priority = bt.max_delimiter_priority(exclude={id(last_leaf)})
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2669
    except ValueError:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2670
        raise CannotSplit("No delimiters found")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2671
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2672
    if delimiter_priority == DOT_PRIORITY:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2673
        if bt.delimiter_count_with_priority(delimiter_priority) == 1:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2674
            raise CannotSplit("Splitting a single attribute from its owner looks wrong")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2675
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2676
    current_line = Line(depth=line.depth, inside_brackets=line.inside_brackets)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2677
    lowest_depth = sys.maxsize
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2678
    trailing_comma_safe = True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2679
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2680
    def append_to_line(leaf: Leaf) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2681
        """Append `leaf` to current line or to new line if appending impossible."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2682
        nonlocal current_line
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2683
        try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2684
            current_line.append_safe(leaf, preformatted=True)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2685
        except ValueError:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2686
            yield current_line
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2687
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2688
            current_line = Line(depth=line.depth, inside_brackets=line.inside_brackets)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2689
            current_line.append(leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2690
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2691
    for leaf in line.leaves:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2692
        yield from append_to_line(leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2693
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2694
        for comment_after in line.comments_after(leaf):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2695
            yield from append_to_line(comment_after)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2696
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2697
        lowest_depth = min(lowest_depth, leaf.bracket_depth)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2698
        if leaf.bracket_depth == lowest_depth:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2699
            if is_vararg(leaf, within={syms.typedargslist}):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2700
                trailing_comma_safe = (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2701
                    trailing_comma_safe and Feature.TRAILING_COMMA_IN_DEF in features
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2702
                )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2703
            elif is_vararg(leaf, within={syms.arglist, syms.argument}):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2704
                trailing_comma_safe = (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2705
                    trailing_comma_safe and Feature.TRAILING_COMMA_IN_CALL in features
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2706
                )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2707
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2708
        leaf_priority = bt.delimiters.get(id(leaf))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2709
        if leaf_priority == delimiter_priority:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2710
            yield current_line
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2711
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2712
            current_line = Line(depth=line.depth, inside_brackets=line.inside_brackets)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2713
    if current_line:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2714
        if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2715
            trailing_comma_safe
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2716
            and delimiter_priority == COMMA_PRIORITY
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2717
            and current_line.leaves[-1].type != token.COMMA
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2718
            and current_line.leaves[-1].type != STANDALONE_COMMENT
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2719
        ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2720
            current_line.append(Leaf(token.COMMA, ","))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2721
        yield current_line
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2722
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2723
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2724
@dont_increase_indentation
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2725
def standalone_comment_split(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2726
    line: Line, features: Collection[Feature] = ()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2727
) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2728
    """Split standalone comments from the rest of the line."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2729
    if not line.contains_standalone_comments(0):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2730
        raise CannotSplit("Line does not have any standalone comments")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2731
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2732
    current_line = Line(depth=line.depth, inside_brackets=line.inside_brackets)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2733
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2734
    def append_to_line(leaf: Leaf) -> Iterator[Line]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2735
        """Append `leaf` to current line or to new line if appending impossible."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2736
        nonlocal current_line
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2737
        try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2738
            current_line.append_safe(leaf, preformatted=True)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2739
        except ValueError:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2740
            yield current_line
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2741
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2742
            current_line = Line(depth=line.depth, inside_brackets=line.inside_brackets)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2743
            current_line.append(leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2744
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2745
    for leaf in line.leaves:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2746
        yield from append_to_line(leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2747
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2748
        for comment_after in line.comments_after(leaf):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2749
            yield from append_to_line(comment_after)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2750
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2751
    if current_line:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2752
        yield current_line
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2753
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2754
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2755
def is_import(leaf: Leaf) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2756
    """Return True if the given leaf starts an import statement."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2757
    p = leaf.parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2758
    t = leaf.type
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2759
    v = leaf.value
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2760
    return bool(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2761
        t == token.NAME
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2762
        and (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2763
            (v == "import" and p and p.type == syms.import_name)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2764
            or (v == "from" and p and p.type == syms.import_from)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2765
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2766
    )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2767
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2768
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2769
def is_type_comment(leaf: Leaf, suffix: str = "") -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2770
    """Return True if the given leaf is a special comment.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2771
    Only returns true for type comments for now."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2772
    t = leaf.type
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2773
    v = leaf.value
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2774
    return t in {token.COMMENT, t == STANDALONE_COMMENT} and v.startswith(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2775
        "# type:" + suffix
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2776
    )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2777
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2778
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2779
def normalize_prefix(leaf: Leaf, *, inside_brackets: bool) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2780
    """Leave existing extra newlines if not `inside_brackets`. Remove everything
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2781
    else.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2782
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2783
    Note: don't use backslashes for formatting or you'll lose your voting rights.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2784
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2785
    if not inside_brackets:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2786
        spl = leaf.prefix.split("#")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2787
        if "\\" not in spl[0]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2788
            nl_count = spl[-1].count("\n")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2789
            if len(spl) > 1:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2790
                nl_count -= 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2791
            leaf.prefix = "\n" * nl_count
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2792
            return
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2793
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2794
    leaf.prefix = ""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2795
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2796
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2797
def normalize_string_prefix(leaf: Leaf, remove_u_prefix: bool = False) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2798
    """Make all string prefixes lowercase.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2799
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2800
    If remove_u_prefix is given, also removes any u prefix from the string.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2801
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2802
    Note: Mutates its argument.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2803
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2804
    match = re.match(r"^([furbFURB]*)(.*)$", leaf.value, re.DOTALL)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2805
    assert match is not None, f"failed to match string {leaf.value!r}"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2806
    orig_prefix = match.group(1)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2807
    new_prefix = orig_prefix.lower()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2808
    if remove_u_prefix:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2809
        new_prefix = new_prefix.replace("u", "")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2810
    leaf.value = f"{new_prefix}{match.group(2)}"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2811
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2812
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2813
def normalize_string_quotes(leaf: Leaf) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2814
    """Prefer double quotes but only if it doesn't cause more escaping.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2815
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2816
    Adds or removes backslashes as appropriate. Doesn't parse and fix
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2817
    strings nested in f-strings (yet).
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2818
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2819
    Note: Mutates its argument.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2820
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2821
    value = leaf.value.lstrip("furbFURB")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2822
    if value[:3] == '"""':
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2823
        return
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2824
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2825
    elif value[:3] == "'''":
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2826
        orig_quote = "'''"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2827
        new_quote = '"""'
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2828
    elif value[0] == '"':
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2829
        orig_quote = '"'
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2830
        new_quote = "'"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2831
    else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2832
        orig_quote = "'"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2833
        new_quote = '"'
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2834
    first_quote_pos = leaf.value.find(orig_quote)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2835
    if first_quote_pos == -1:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2836
        return  # There's an internal error
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2837
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2838
    prefix = leaf.value[:first_quote_pos]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2839
    unescaped_new_quote = re.compile(rf"(([^\\]|^)(\\\\)*){new_quote}")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2840
    escaped_new_quote = re.compile(rf"([^\\]|^)\\((?:\\\\)*){new_quote}")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2841
    escaped_orig_quote = re.compile(rf"([^\\]|^)\\((?:\\\\)*){orig_quote}")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2842
    body = leaf.value[first_quote_pos + len(orig_quote) : -len(orig_quote)]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2843
    if "r" in prefix.casefold():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2844
        if unescaped_new_quote.search(body):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2845
            # There's at least one unescaped new_quote in this raw string
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2846
            # so converting is impossible
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2847
            return
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2848
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2849
        # Do not introduce or remove backslashes in raw strings
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2850
        new_body = body
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2851
    else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2852
        # remove unnecessary escapes
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2853
        new_body = sub_twice(escaped_new_quote, rf"\1\2{new_quote}", body)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2854
        if body != new_body:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2855
            # Consider the string without unnecessary escapes as the original
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2856
            body = new_body
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2857
            leaf.value = f"{prefix}{orig_quote}{body}{orig_quote}"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2858
        new_body = sub_twice(escaped_orig_quote, rf"\1\2{orig_quote}", new_body)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2859
        new_body = sub_twice(unescaped_new_quote, rf"\1\\{new_quote}", new_body)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2860
    if "f" in prefix.casefold():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2861
        matches = re.findall(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2862
            r"""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2863
            (?:[^{]|^)\{  # start of the string or a non-{ followed by a single {
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2864
                ([^{].*?)  # contents of the brackets except if begins with {{
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2865
            \}(?:[^}]|$)  # A } followed by end of the string or a non-}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2866
            """,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2867
            new_body,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2868
            re.VERBOSE,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2869
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2870
        for m in matches:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2871
            if "\\" in str(m):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2872
                # Do not introduce backslashes in interpolated expressions
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2873
                return
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2874
    if new_quote == '"""' and new_body[-1:] == '"':
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2875
        # edge case:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2876
        new_body = new_body[:-1] + '\\"'
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2877
    orig_escape_count = body.count("\\")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2878
    new_escape_count = new_body.count("\\")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2879
    if new_escape_count > orig_escape_count:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2880
        return  # Do not introduce more escaping
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2881
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2882
    if new_escape_count == orig_escape_count and orig_quote == '"':
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2883
        return  # Prefer double quotes
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2884
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2885
    leaf.value = f"{prefix}{new_quote}{new_body}{new_quote}"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2886
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2887
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2888
def normalize_numeric_literal(leaf: Leaf) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2889
    """Normalizes numeric (float, int, and complex) literals.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2890
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2891
    All letters used in the representation are normalized to lowercase (except
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2892
    in Python 2 long literals).
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2893
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2894
    text = leaf.value.lower()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2895
    if text.startswith(("0o", "0b")):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2896
        # Leave octal and binary literals alone.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2897
        pass
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2898
    elif text.startswith("0x"):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2899
        # Change hex literals to upper case.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2900
        before, after = text[:2], text[2:]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2901
        text = f"{before}{after.upper()}"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2902
    elif "e" in text:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2903
        before, after = text.split("e")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2904
        sign = ""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2905
        if after.startswith("-"):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2906
            after = after[1:]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2907
            sign = "-"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2908
        elif after.startswith("+"):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2909
            after = after[1:]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2910
        before = format_float_or_int_string(before)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2911
        text = f"{before}e{sign}{after}"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2912
    elif text.endswith(("j", "l")):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2913
        number = text[:-1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2914
        suffix = text[-1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2915
        # Capitalize in "2L" because "l" looks too similar to "1".
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2916
        if suffix == "l":
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2917
            suffix = "L"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2918
        text = f"{format_float_or_int_string(number)}{suffix}"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2919
    else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2920
        text = format_float_or_int_string(text)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2921
    leaf.value = text
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2922
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2923
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2924
def format_float_or_int_string(text: str) -> str:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2925
    """Formats a float string like "1.0"."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2926
    if "." not in text:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2927
        return text
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2928
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2929
    before, after = text.split(".")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2930
    return f"{before or 0}.{after or 0}"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2931
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2932
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2933
def normalize_invisible_parens(node: Node, parens_after: Set[str]) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2934
    """Make existing optional parentheses invisible or create new ones.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2935
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2936
    `parens_after` is a set of string leaf values immediately after which parens
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2937
    should be put.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2938
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2939
    Standardizes on visible parentheses for single-element tuples, and keeps
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2940
    existing visible parentheses for other tuples and generator expressions.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2941
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2942
    for pc in list_comments(node.prefix, is_endmarker=False):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2943
        if pc.value in FMT_OFF:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2944
            # This `node` has a prefix with `# fmt: off`, don't mess with parens.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2945
            return
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2946
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2947
    check_lpar = False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2948
    for index, child in enumerate(list(node.children)):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2949
        # Add parentheses around long tuple unpacking in assignments.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2950
        if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2951
            index == 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2952
            and isinstance(child, Node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2953
            and child.type == syms.testlist_star_expr
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2954
        ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2955
            check_lpar = True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2956
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2957
        if check_lpar:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2958
            if is_walrus_assignment(child):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2959
                continue
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2960
            if child.type == syms.atom:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2961
                # Determines if the underlying atom should be surrounded with
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2962
                # invisible params - also makes parens invisible recursively
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2963
                # within the atom and removes repeated invisible parens within
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2964
                # the atom
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2965
                should_surround_with_parens = maybe_make_parens_invisible_in_atom(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2966
                    child, parent=node
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2967
                )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2968
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2969
                if should_surround_with_parens:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2970
                    lpar = Leaf(token.LPAR, "")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2971
                    rpar = Leaf(token.RPAR, "")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2972
                    index = child.remove() or 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2973
                    node.insert_child(index, Node(syms.atom, [lpar, child, rpar]))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2974
            elif is_one_tuple(child):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2975
                # wrap child in visible parentheses
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2976
                lpar = Leaf(token.LPAR, "(")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2977
                rpar = Leaf(token.RPAR, ")")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2978
                child.remove()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2979
                node.insert_child(index, Node(syms.atom, [lpar, child, rpar]))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2980
            elif node.type == syms.import_from:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2981
                # "import from" nodes store parentheses directly as part of
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2982
                # the statement
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2983
                if child.type == token.LPAR:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2984
                    # make parentheses invisible
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2985
                    child.value = ""  # type: ignore
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2986
                    node.children[-1].value = ""  # type: ignore
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2987
                elif child.type != token.STAR:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2988
                    # insert invisible parentheses
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2989
                    node.insert_child(index, Leaf(token.LPAR, ""))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2990
                    node.append_child(Leaf(token.RPAR, ""))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2991
                break
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2992
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2993
            elif not (isinstance(child, Leaf) and is_multiline_string(child)):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2994
                # wrap child in invisible parentheses
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2995
                lpar = Leaf(token.LPAR, "")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2996
                rpar = Leaf(token.RPAR, "")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2997
                index = child.remove() or 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2998
                prefix = child.prefix
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  2999
                child.prefix = ""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3000
                new_child = Node(syms.atom, [lpar, child, rpar])
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3001
                new_child.prefix = prefix
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3002
                node.insert_child(index, new_child)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3003
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3004
        check_lpar = isinstance(child, Leaf) and child.value in parens_after
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3005
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3006
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3007
def normalize_fmt_off(node: Node) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3008
    """Convert content between `# fmt: off`/`# fmt: on` into standalone comments."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3009
    try_again = True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3010
    while try_again:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3011
        try_again = convert_one_fmt_off_pair(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3012
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3013
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3014
def convert_one_fmt_off_pair(node: Node) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3015
    """Convert content of a single `# fmt: off`/`# fmt: on` into a standalone comment.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3016
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3017
    Returns True if a pair was converted.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3018
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3019
    for leaf in node.leaves():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3020
        previous_consumed = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3021
        for comment in list_comments(leaf.prefix, is_endmarker=False):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3022
            if comment.value in FMT_OFF:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3023
                # We only want standalone comments. If there's no previous leaf or
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3024
                # the previous leaf is indentation, it's a standalone comment in
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3025
                # disguise.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3026
                if comment.type != STANDALONE_COMMENT:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3027
                    prev = preceding_leaf(leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3028
                    if prev and prev.type not in WHITESPACE:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3029
                        continue
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3030
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3031
                ignored_nodes = list(generate_ignored_nodes(leaf))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3032
                if not ignored_nodes:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3033
                    continue
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3034
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3035
                first = ignored_nodes[0]  # Can be a container node with the `leaf`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3036
                parent = first.parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3037
                prefix = first.prefix
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3038
                first.prefix = prefix[comment.consumed :]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3039
                hidden_value = (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3040
                    comment.value + "\n" + "".join(str(n) for n in ignored_nodes)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3041
                )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3042
                if hidden_value.endswith("\n"):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3043
                    # That happens when one of the `ignored_nodes` ended with a NEWLINE
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3044
                    # leaf (possibly followed by a DEDENT).
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3045
                    hidden_value = hidden_value[:-1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3046
                first_idx = None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3047
                for ignored in ignored_nodes:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3048
                    index = ignored.remove()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3049
                    if first_idx is None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3050
                        first_idx = index
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3051
                assert parent is not None, "INTERNAL ERROR: fmt: on/off handling (1)"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3052
                assert first_idx is not None, "INTERNAL ERROR: fmt: on/off handling (2)"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3053
                parent.insert_child(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3054
                    first_idx,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3055
                    Leaf(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3056
                        STANDALONE_COMMENT,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3057
                        hidden_value,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3058
                        prefix=prefix[:previous_consumed] + "\n" * comment.newlines,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3059
                    ),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3060
                )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3061
                return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3062
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3063
            previous_consumed = comment.consumed
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3064
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3065
    return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3066
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3067
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3068
def generate_ignored_nodes(leaf: Leaf) -> Iterator[LN]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3069
    """Starting from the container of `leaf`, generate all leaves until `# fmt: on`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3070
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3071
    Stops at the end of the block.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3072
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3073
    container: Optional[LN] = container_of(leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3074
    while container is not None and container.type != token.ENDMARKER:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3075
        for comment in list_comments(container.prefix, is_endmarker=False):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3076
            if comment.value in FMT_ON:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3077
                return
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3078
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3079
        yield container
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3080
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3081
        container = container.next_sibling
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3082
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3083
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3084
def maybe_make_parens_invisible_in_atom(node: LN, parent: LN) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3085
    """If it's safe, make the parens in the atom `node` invisible, recursively.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3086
    Additionally, remove repeated, adjacent invisible parens from the atom `node`
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3087
    as they are redundant.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3088
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3089
    Returns whether the node should itself be wrapped in invisible parentheses.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3090
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3091
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3092
    if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3093
        node.type != syms.atom
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3094
        or is_empty_tuple(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3095
        or is_one_tuple(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3096
        or (is_yield(node) and parent.type != syms.expr_stmt)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3097
        or max_delimiter_priority_in_atom(node) >= COMMA_PRIORITY
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3098
    ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3099
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3100
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3101
    first = node.children[0]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3102
    last = node.children[-1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3103
    if first.type == token.LPAR and last.type == token.RPAR:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3104
        middle = node.children[1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3105
        # make parentheses invisible
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3106
        first.value = ""  # type: ignore
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3107
        last.value = ""  # type: ignore
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3108
        maybe_make_parens_invisible_in_atom(middle, parent=parent)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3109
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3110
        if is_atom_with_invisible_parens(middle):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3111
            # Strip the invisible parens from `middle` by replacing
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3112
            # it with the child in-between the invisible parens
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3113
            middle.replace(middle.children[1])
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3114
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3115
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3116
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3117
    return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3118
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3119
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3120
def is_atom_with_invisible_parens(node: LN) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3121
    """Given a `LN`, determines whether it's an atom `node` with invisible
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3122
    parens. Useful in dedupe-ing and normalizing parens.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3123
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3124
    if isinstance(node, Leaf) or node.type != syms.atom:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3125
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3126
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3127
    first, last = node.children[0], node.children[-1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3128
    return (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3129
        isinstance(first, Leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3130
        and first.type == token.LPAR
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3131
        and first.value == ""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3132
        and isinstance(last, Leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3133
        and last.type == token.RPAR
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3134
        and last.value == ""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3135
    )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3136
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3137
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3138
def is_empty_tuple(node: LN) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3139
    """Return True if `node` holds an empty tuple."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3140
    return (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3141
        node.type == syms.atom
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3142
        and len(node.children) == 2
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3143
        and node.children[0].type == token.LPAR
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3144
        and node.children[1].type == token.RPAR
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3145
    )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3146
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3147
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3148
def unwrap_singleton_parenthesis(node: LN) -> Optional[LN]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3149
    """Returns `wrapped` if `node` is of the shape ( wrapped ).
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3150
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3151
    Parenthesis can be optional. Returns None otherwise"""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3152
    if len(node.children) != 3:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3153
        return None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3154
    lpar, wrapped, rpar = node.children
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3155
    if not (lpar.type == token.LPAR and rpar.type == token.RPAR):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3156
        return None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3157
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3158
    return wrapped
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3159
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3160
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3161
def is_one_tuple(node: LN) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3162
    """Return True if `node` holds a tuple with one element, with or without parens."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3163
    if node.type == syms.atom:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3164
        gexp = unwrap_singleton_parenthesis(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3165
        if gexp is None or gexp.type != syms.testlist_gexp:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3166
            return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3167
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3168
        return len(gexp.children) == 2 and gexp.children[1].type == token.COMMA
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3169
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3170
    return (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3171
        node.type in IMPLICIT_TUPLE
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3172
        and len(node.children) == 2
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3173
        and node.children[1].type == token.COMMA
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3174
    )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3175
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3176
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3177
def is_walrus_assignment(node: LN) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3178
    """Return True iff `node` is of the shape ( test := test )"""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3179
    inner = unwrap_singleton_parenthesis(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3180
    return inner is not None and inner.type == syms.namedexpr_test
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3181
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3182
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3183
def is_yield(node: LN) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3184
    """Return True if `node` holds a `yield` or `yield from` expression."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3185
    if node.type == syms.yield_expr:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3186
        return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3187
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3188
    if node.type == token.NAME and node.value == "yield":  # type: ignore
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3189
        return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3190
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3191
    if node.type != syms.atom:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3192
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3193
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3194
    if len(node.children) != 3:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3195
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3196
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3197
    lpar, expr, rpar = node.children
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3198
    if lpar.type == token.LPAR and rpar.type == token.RPAR:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3199
        return is_yield(expr)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3200
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3201
    return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3202
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3203
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3204
def is_vararg(leaf: Leaf, within: Set[NodeType]) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3205
    """Return True if `leaf` is a star or double star in a vararg or kwarg.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3206
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3207
    If `within` includes VARARGS_PARENTS, this applies to function signatures.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3208
    If `within` includes UNPACKING_PARENTS, it applies to right hand-side
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3209
    extended iterable unpacking (PEP 3132) and additional unpacking
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3210
    generalizations (PEP 448).
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3211
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3212
    if leaf.type not in VARARGS_SPECIALS or not leaf.parent:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3213
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3214
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3215
    p = leaf.parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3216
    if p.type == syms.star_expr:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3217
        # Star expressions are also used as assignment targets in extended
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3218
        # iterable unpacking (PEP 3132).  See what its parent is instead.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3219
        if not p.parent:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3220
            return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3221
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3222
        p = p.parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3223
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3224
    return p.type in within
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3225
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3226
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3227
def is_multiline_string(leaf: Leaf) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3228
    """Return True if `leaf` is a multiline string that actually spans many lines."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3229
    value = leaf.value.lstrip("furbFURB")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3230
    return value[:3] in {'"""', "'''"} and "\n" in value
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3231
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3232
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3233
def is_stub_suite(node: Node) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3234
    """Return True if `node` is a suite with a stub body."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3235
    if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3236
        len(node.children) != 4
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3237
        or node.children[0].type != token.NEWLINE
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3238
        or node.children[1].type != token.INDENT
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3239
        or node.children[3].type != token.DEDENT
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3240
    ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3241
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3242
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3243
    return is_stub_body(node.children[2])
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3244
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3245
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3246
def is_stub_body(node: LN) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3247
    """Return True if `node` is a simple statement containing an ellipsis."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3248
    if not isinstance(node, Node) or node.type != syms.simple_stmt:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3249
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3250
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3251
    if len(node.children) != 2:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3252
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3253
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3254
    child = node.children[0]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3255
    return (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3256
        child.type == syms.atom
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3257
        and len(child.children) == 3
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3258
        and all(leaf == Leaf(token.DOT, ".") for leaf in child.children)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3259
    )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3260
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3261
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3262
def max_delimiter_priority_in_atom(node: LN) -> Priority:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3263
    """Return maximum delimiter priority inside `node`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3264
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3265
    This is specific to atoms with contents contained in a pair of parentheses.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3266
    If `node` isn't an atom or there are no enclosing parentheses, returns 0.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3267
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3268
    if node.type != syms.atom:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3269
        return 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3270
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3271
    first = node.children[0]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3272
    last = node.children[-1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3273
    if not (first.type == token.LPAR and last.type == token.RPAR):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3274
        return 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3275
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3276
    bt = BracketTracker()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3277
    for c in node.children[1:-1]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3278
        if isinstance(c, Leaf):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3279
            bt.mark(c)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3280
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3281
            for leaf in c.leaves():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3282
                bt.mark(leaf)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3283
    try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3284
        return bt.max_delimiter_priority()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3285
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3286
    except ValueError:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3287
        return 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3288
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3289
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3290
def ensure_visible(leaf: Leaf) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3291
    """Make sure parentheses are visible.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3292
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3293
    They could be invisible as part of some statements (see
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3294
    :func:`normalize_invisible_parens` and :func:`visit_import_from`).
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3295
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3296
    if leaf.type == token.LPAR:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3297
        leaf.value = "("
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3298
    elif leaf.type == token.RPAR:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3299
        leaf.value = ")"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3300
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3301
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3302
def should_explode(line: Line, opening_bracket: Leaf) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3303
    """Should `line` immediately be split with `delimiter_split()` after RHS?"""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3304
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3305
    if not (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3306
        opening_bracket.parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3307
        and opening_bracket.parent.type in {syms.atom, syms.import_from}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3308
        and opening_bracket.value in "[{("
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3309
    ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3310
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3311
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3312
    try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3313
        last_leaf = line.leaves[-1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3314
        exclude = {id(last_leaf)} if last_leaf.type == token.COMMA else set()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3315
        max_priority = line.bracket_tracker.max_delimiter_priority(exclude=exclude)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3316
    except (IndexError, ValueError):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3317
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3318
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3319
    return max_priority == COMMA_PRIORITY
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3320
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3321
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3322
def get_features_used(node: Node) -> Set[Feature]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3323
    """Return a set of (relatively) new Python features used in this file.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3324
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3325
    Currently looking for:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3326
    - f-strings;
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3327
    - underscores in numeric literals;
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3328
    - trailing commas after * or ** in function signatures and calls;
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3329
    - positional only arguments in function signatures and lambdas;
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3330
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3331
    features: Set[Feature] = set()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3332
    for n in node.pre_order():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3333
        if n.type == token.STRING:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3334
            value_head = n.value[:2]  # type: ignore
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3335
            if value_head in {'f"', 'F"', "f'", "F'", "rf", "fr", "RF", "FR"}:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3336
                features.add(Feature.F_STRINGS)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3337
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3338
        elif n.type == token.NUMBER:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3339
            if "_" in n.value:  # type: ignore
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3340
                features.add(Feature.NUMERIC_UNDERSCORES)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3341
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3342
        elif n.type == token.SLASH:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3343
            if n.parent and n.parent.type in {syms.typedargslist, syms.arglist}:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3344
                features.add(Feature.POS_ONLY_ARGUMENTS)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3345
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3346
        elif n.type == token.COLONEQUAL:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3347
            features.add(Feature.ASSIGNMENT_EXPRESSIONS)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3348
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3349
        elif (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3350
            n.type in {syms.typedargslist, syms.arglist}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3351
            and n.children
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3352
            and n.children[-1].type == token.COMMA
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3353
        ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3354
            if n.type == syms.typedargslist:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3355
                feature = Feature.TRAILING_COMMA_IN_DEF
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3356
            else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3357
                feature = Feature.TRAILING_COMMA_IN_CALL
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3358
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3359
            for ch in n.children:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3360
                if ch.type in STARS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3361
                    features.add(feature)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3362
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3363
                if ch.type == syms.argument:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3364
                    for argch in ch.children:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3365
                        if argch.type in STARS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3366
                            features.add(feature)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3367
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3368
    return features
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3369
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3370
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3371
def detect_target_versions(node: Node) -> Set[TargetVersion]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3372
    """Detect the version to target based on the nodes used."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3373
    features = get_features_used(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3374
    return {
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3375
        version for version in TargetVersion if features <= VERSION_TO_FEATURES[version]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3376
    }
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3377
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3378
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3379
def generate_trailers_to_omit(line: Line, line_length: int) -> Iterator[Set[LeafID]]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3380
    """Generate sets of closing bracket IDs that should be omitted in a RHS.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3381
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3382
    Brackets can be omitted if the entire trailer up to and including
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3383
    a preceding closing bracket fits in one line.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3384
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3385
    Yielded sets are cumulative (contain results of previous yields, too).  First
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3386
    set is empty.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3387
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3388
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3389
    omit: Set[LeafID] = set()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3390
    yield omit
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3391
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3392
    length = 4 * line.depth
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3393
    opening_bracket = None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3394
    closing_bracket = None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3395
    inner_brackets: Set[LeafID] = set()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3396
    for index, leaf, leaf_length in enumerate_with_length(line, reversed=True):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3397
        length += leaf_length
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3398
        if length > line_length:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3399
            break
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3400
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3401
        has_inline_comment = leaf_length > len(leaf.value) + len(leaf.prefix)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3402
        if leaf.type == STANDALONE_COMMENT or has_inline_comment:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3403
            break
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3404
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3405
        if opening_bracket:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3406
            if leaf is opening_bracket:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3407
                opening_bracket = None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3408
            elif leaf.type in CLOSING_BRACKETS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3409
                inner_brackets.add(id(leaf))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3410
        elif leaf.type in CLOSING_BRACKETS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3411
            if index > 0 and line.leaves[index - 1].type in OPENING_BRACKETS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3412
                # Empty brackets would fail a split so treat them as "inner"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3413
                # brackets (e.g. only add them to the `omit` set if another
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3414
                # pair of brackets was good enough.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3415
                inner_brackets.add(id(leaf))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3416
                continue
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3417
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3418
            if closing_bracket:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3419
                omit.add(id(closing_bracket))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3420
                omit.update(inner_brackets)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3421
                inner_brackets.clear()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3422
                yield omit
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3423
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3424
            if leaf.value:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3425
                opening_bracket = leaf.opening_bracket
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3426
                closing_bracket = leaf
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3427
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3428
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3429
def get_future_imports(node: Node) -> Set[str]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3430
    """Return a set of __future__ imports in the file."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3431
    imports: Set[str] = set()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3432
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3433
    def get_imports_from_children(children: List[LN]) -> Generator[str, None, None]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3434
        for child in children:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3435
            if isinstance(child, Leaf):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3436
                if child.type == token.NAME:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3437
                    yield child.value
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3438
            elif child.type == syms.import_as_name:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3439
                orig_name = child.children[0]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3440
                assert isinstance(orig_name, Leaf), "Invalid syntax parsing imports"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3441
                assert orig_name.type == token.NAME, "Invalid syntax parsing imports"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3442
                yield orig_name.value
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3443
            elif child.type == syms.import_as_names:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3444
                yield from get_imports_from_children(child.children)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3445
            else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3446
                raise AssertionError("Invalid syntax parsing imports")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3447
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3448
    for child in node.children:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3449
        if child.type != syms.simple_stmt:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3450
            break
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3451
        first_child = child.children[0]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3452
        if isinstance(first_child, Leaf):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3453
            # Continue looking if we see a docstring; otherwise stop.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3454
            if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3455
                len(child.children) == 2
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3456
                and first_child.type == token.STRING
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3457
                and child.children[1].type == token.NEWLINE
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3458
            ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3459
                continue
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3460
            else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3461
                break
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3462
        elif first_child.type == syms.import_from:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3463
            module_name = first_child.children[1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3464
            if not isinstance(module_name, Leaf) or module_name.value != "__future__":
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3465
                break
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3466
            imports |= set(get_imports_from_children(first_child.children[3:]))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3467
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3468
            break
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3469
    return imports
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3470
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3471
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3472
def gen_python_files_in_dir(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3473
    path: Path,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3474
    root: Path,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3475
    include: Pattern[str],
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3476
    exclude: Pattern[str],
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3477
    report: "Report",
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3478
) -> Iterator[Path]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3479
    """Generate all files under `path` whose paths are not excluded by the
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3480
    `exclude` regex, but are included by the `include` regex.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3481
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3482
    Symbolic links pointing outside of the `root` directory are ignored.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3483
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3484
    `report` is where output about exclusions goes.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3485
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3486
    assert root.is_absolute(), f"INTERNAL ERROR: `root` must be absolute but is {root}"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3487
    for child in path.iterdir():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3488
        try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3489
            normalized_path = "/" + child.resolve().relative_to(root).as_posix()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3490
        except ValueError:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3491
            if child.is_symlink():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3492
                report.path_ignored(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3493
                    child, f"is a symbolic link that points outside {root}"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3494
                )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3495
                continue
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3496
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3497
            raise
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3498
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3499
        if child.is_dir():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3500
            normalized_path += "/"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3501
        exclude_match = exclude.search(normalized_path)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3502
        if exclude_match and exclude_match.group(0):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3503
            report.path_ignored(child, f"matches the --exclude regular expression")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3504
            continue
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3505
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3506
        if child.is_dir():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3507
            yield from gen_python_files_in_dir(child, root, include, exclude, report)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3508
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3509
        elif child.is_file():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3510
            include_match = include.search(normalized_path)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3511
            if include_match:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3512
                yield child
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3513
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3514
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3515
@lru_cache()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3516
def find_project_root(srcs: Iterable[str]) -> Path:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3517
    """Return a directory containing .git, .hg, or pyproject.toml.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3518
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3519
    That directory can be one of the directories passed in `srcs` or their
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3520
    common parent.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3521
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3522
    If no directory in the tree contains a marker that would specify it's the
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3523
    project root, the root of the file system is returned.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3524
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3525
    if not srcs:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3526
        return Path("/").resolve()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3527
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3528
    common_base = min(Path(src).resolve() for src in srcs)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3529
    if common_base.is_dir():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3530
        # Append a fake file so `parents` below returns `common_base_dir`, too.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3531
        common_base /= "fake-file"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3532
    for directory in common_base.parents:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3533
        if (directory / ".git").is_dir():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3534
            return directory
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3535
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3536
        if (directory / ".hg").is_dir():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3537
            return directory
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3538
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3539
        if (directory / "pyproject.toml").is_file():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3540
            return directory
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3541
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3542
    return directory
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3543
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3544
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3545
@dataclass
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3546
class Report:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3547
    """Provides a reformatting counter. Can be rendered with `str(report)`."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3548
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3549
    check: bool = False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3550
    quiet: bool = False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3551
    verbose: bool = False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3552
    change_count: int = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3553
    same_count: int = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3554
    failure_count: int = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3555
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3556
    def done(self, src: Path, changed: Changed) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3557
        """Increment the counter for successful reformatting. Write out a message."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3558
        if changed is Changed.YES:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3559
            reformatted = "would reformat" if self.check else "reformatted"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3560
            if self.verbose or not self.quiet:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3561
                out(f"{reformatted} {src}")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3562
            self.change_count += 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3563
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3564
            if self.verbose:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3565
                if changed is Changed.NO:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3566
                    msg = f"{src} already well formatted, good job."
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3567
                else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3568
                    msg = f"{src} wasn't modified on disk since last run."
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3569
                out(msg, bold=False)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3570
            self.same_count += 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3571
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3572
    def failed(self, src: Path, message: str) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3573
        """Increment the counter for failed reformatting. Write out a message."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3574
        err(f"error: cannot format {src}: {message}")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3575
        self.failure_count += 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3576
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3577
    def path_ignored(self, path: Path, message: str) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3578
        if self.verbose:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3579
            out(f"{path} ignored: {message}", bold=False)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3580
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3581
    @property
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3582
    def return_code(self) -> int:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3583
        """Return the exit code that the app should use.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3584
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3585
        This considers the current state of changed files and failures:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3586
        - if there were any failures, return 123;
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3587
        - if any files were changed and --check is being used, return 1;
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3588
        - otherwise return 0.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3589
        """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3590
        # According to http://tldp.org/LDP/abs/html/exitcodes.html starting with
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3591
        # 126 we have special return codes reserved by the shell.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3592
        if self.failure_count:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3593
            return 123
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3594
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3595
        elif self.change_count and self.check:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3596
            return 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3597
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3598
        return 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3599
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3600
    def __str__(self) -> str:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3601
        """Render a color report of the current state.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3602
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3603
        Use `click.unstyle` to remove colors.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3604
        """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3605
        if self.check:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3606
            reformatted = "would be reformatted"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3607
            unchanged = "would be left unchanged"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3608
            failed = "would fail to reformat"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3609
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3610
            reformatted = "reformatted"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3611
            unchanged = "left unchanged"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3612
            failed = "failed to reformat"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3613
        report = []
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3614
        if self.change_count:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3615
            s = "s" if self.change_count > 1 else ""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3616
            report.append(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3617
                click.style(f"{self.change_count} file{s} {reformatted}", bold=True)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3618
            )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3619
        if self.same_count:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3620
            s = "s" if self.same_count > 1 else ""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3621
            report.append(f"{self.same_count} file{s} {unchanged}")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3622
        if self.failure_count:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3623
            s = "s" if self.failure_count > 1 else ""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3624
            report.append(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3625
                click.style(f"{self.failure_count} file{s} {failed}", fg="red")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3626
            )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3627
        return ", ".join(report) + "."
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3628
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3629
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3630
def parse_ast(src: str) -> Union[ast.AST, ast3.AST, ast27.AST]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3631
    filename = "<unknown>"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3632
    if sys.version_info >= (3, 8):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3633
        # TODO: support Python 4+ ;)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3634
        for minor_version in range(sys.version_info[1], 4, -1):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3635
            try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3636
                return ast.parse(src, filename, feature_version=(3, minor_version))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3637
            except SyntaxError:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3638
                continue
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3639
    else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3640
        for feature_version in (7, 6):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3641
            try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3642
                return ast3.parse(src, filename, feature_version=feature_version)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3643
            except SyntaxError:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3644
                continue
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3645
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3646
    return ast27.parse(src)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3647
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3648
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3649
def _fixup_ast_constants(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3650
    node: Union[ast.AST, ast3.AST, ast27.AST]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3651
) -> Union[ast.AST, ast3.AST, ast27.AST]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3652
    """Map ast nodes deprecated in 3.8 to Constant."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3653
    # casts are required until this is released:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3654
    # https://github.com/python/typeshed/pull/3142
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3655
    if isinstance(node, (ast.Str, ast3.Str, ast27.Str, ast.Bytes, ast3.Bytes)):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3656
        return cast(ast.AST, ast.Constant(value=node.s))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3657
    elif isinstance(node, (ast.Num, ast3.Num, ast27.Num)):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3658
        return cast(ast.AST, ast.Constant(value=node.n))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3659
    elif isinstance(node, (ast.NameConstant, ast3.NameConstant)):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3660
        return cast(ast.AST, ast.Constant(value=node.value))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3661
    return node
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3662
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3663
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3664
def assert_equivalent(src: str, dst: str) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3665
    """Raise AssertionError if `src` and `dst` aren't equivalent."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3666
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3667
    def _v(node: Union[ast.AST, ast3.AST, ast27.AST], depth: int = 0) -> Iterator[str]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3668
        """Simple visitor generating strings to compare ASTs by content."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3669
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3670
        node = _fixup_ast_constants(node)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3671
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3672
        yield f"{'  ' * depth}{node.__class__.__name__}("
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3673
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3674
        for field in sorted(node._fields):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3675
            # TypeIgnore has only one field 'lineno' which breaks this comparison
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3676
            type_ignore_classes = (ast3.TypeIgnore, ast27.TypeIgnore)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3677
            if sys.version_info >= (3, 8):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3678
                type_ignore_classes += (ast.TypeIgnore,)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3679
            if isinstance(node, type_ignore_classes):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3680
                break
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3681
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3682
            try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3683
                value = getattr(node, field)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3684
            except AttributeError:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3685
                continue
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3686
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3687
            yield f"{'  ' * (depth+1)}{field}="
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3688
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3689
            if isinstance(value, list):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3690
                for item in value:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3691
                    # Ignore nested tuples within del statements, because we may insert
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3692
                    # parentheses and they change the AST.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3693
                    if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3694
                        field == "targets"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3695
                        and isinstance(node, (ast.Delete, ast3.Delete, ast27.Delete))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3696
                        and isinstance(item, (ast.Tuple, ast3.Tuple, ast27.Tuple))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3697
                    ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3698
                        for item in item.elts:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3699
                            yield from _v(item, depth + 2)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3700
                    elif isinstance(item, (ast.AST, ast3.AST, ast27.AST)):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3701
                        yield from _v(item, depth + 2)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3702
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3703
            elif isinstance(value, (ast.AST, ast3.AST, ast27.AST)):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3704
                yield from _v(value, depth + 2)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3705
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3706
            else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3707
                yield f"{'  ' * (depth+2)}{value!r},  # {value.__class__.__name__}"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3708
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3709
        yield f"{'  ' * depth})  # /{node.__class__.__name__}"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3710
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3711
    try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3712
        src_ast = parse_ast(src)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3713
    except Exception as exc:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3714
        raise AssertionError(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3715
            f"cannot use --safe with this file; failed to parse source file.  "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3716
            f"AST error message: {exc}"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3717
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3718
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3719
    try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3720
        dst_ast = parse_ast(dst)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3721
    except Exception as exc:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3722
        log = dump_to_file("".join(traceback.format_tb(exc.__traceback__)), dst)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3723
        raise AssertionError(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3724
            f"INTERNAL ERROR: Black produced invalid code: {exc}. "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3725
            f"Please report a bug on https://github.com/psf/black/issues.  "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3726
            f"This invalid output might be helpful: {log}"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3727
        ) from None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3728
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3729
    src_ast_str = "\n".join(_v(src_ast))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3730
    dst_ast_str = "\n".join(_v(dst_ast))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3731
    if src_ast_str != dst_ast_str:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3732
        log = dump_to_file(diff(src_ast_str, dst_ast_str, "src", "dst"))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3733
        raise AssertionError(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3734
            f"INTERNAL ERROR: Black produced code that is not equivalent to "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3735
            f"the source.  "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3736
            f"Please report a bug on https://github.com/psf/black/issues.  "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3737
            f"This diff might be helpful: {log}"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3738
        ) from None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3739
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3740
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3741
def assert_stable(src: str, dst: str, mode: FileMode) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3742
    """Raise AssertionError if `dst` reformats differently the second time."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3743
    newdst = format_str(dst, mode=mode)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3744
    if dst != newdst:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3745
        log = dump_to_file(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3746
            diff(src, dst, "source", "first pass"),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3747
            diff(dst, newdst, "first pass", "second pass"),
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3748
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3749
        raise AssertionError(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3750
            f"INTERNAL ERROR: Black produced different code on the second pass "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3751
            f"of the formatter.  "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3752
            f"Please report a bug on https://github.com/psf/black/issues.  "
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3753
            f"This diff might be helpful: {log}"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3754
        ) from None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3755
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3756
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3757
def dump_to_file(*output: str) -> str:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3758
    """Dump `output` to a temporary file. Return path to the file."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3759
    with tempfile.NamedTemporaryFile(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3760
        mode="w", prefix="blk_", suffix=".log", delete=False, encoding="utf8"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3761
    ) as f:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3762
        for lines in output:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3763
            f.write(lines)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3764
            if lines and lines[-1] != "\n":
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3765
                f.write("\n")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3766
    return f.name
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3767
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3768
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3769
@contextmanager
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3770
def nullcontext() -> Iterator[None]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3771
    """Return context manager that does nothing.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3772
    Similar to `nullcontext` from python 3.7"""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3773
    yield
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3774
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3775
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3776
def diff(a: str, b: str, a_name: str, b_name: str) -> str:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3777
    """Return a unified diff string between strings `a` and `b`."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3778
    import difflib
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3779
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3780
    a_lines = [line + "\n" for line in a.split("\n")]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3781
    b_lines = [line + "\n" for line in b.split("\n")]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3782
    return "".join(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3783
        difflib.unified_diff(a_lines, b_lines, fromfile=a_name, tofile=b_name, n=5)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3784
    )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3785
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3786
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3787
def cancel(tasks: Iterable[asyncio.Task]) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3788
    """asyncio signal handler that cancels all `tasks` and reports to stderr."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3789
    err("Aborted!")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3790
    for task in tasks:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3791
        task.cancel()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3792
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3793
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3794
def shutdown(loop: asyncio.AbstractEventLoop) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3795
    """Cancel all pending tasks on `loop`, wait for them, and close the loop."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3796
    try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3797
        if sys.version_info[:2] >= (3, 7):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3798
            all_tasks = asyncio.all_tasks
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3799
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3800
            all_tasks = asyncio.Task.all_tasks
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3801
        # This part is borrowed from asyncio/runners.py in Python 3.7b2.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3802
        to_cancel = [task for task in all_tasks(loop) if not task.done()]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3803
        if not to_cancel:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3804
            return
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3805
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3806
        for task in to_cancel:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3807
            task.cancel()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3808
        loop.run_until_complete(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3809
            asyncio.gather(*to_cancel, loop=loop, return_exceptions=True)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3810
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3811
    finally:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3812
        # `concurrent.futures.Future` objects cannot be cancelled once they
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3813
        # are already running. There might be some when the `shutdown()` happened.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3814
        # Silence their logger's spew about the event loop being closed.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3815
        cf_logger = logging.getLogger("concurrent.futures")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3816
        cf_logger.setLevel(logging.CRITICAL)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3817
        loop.close()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3818
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3819
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3820
def sub_twice(regex: Pattern[str], replacement: str, original: str) -> str:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3821
    """Replace `regex` with `replacement` twice on `original`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3822
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3823
    This is used by string normalization to perform replaces on
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3824
    overlapping matches.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3825
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3826
    return regex.sub(replacement, regex.sub(replacement, original))
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3827
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3828
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3829
def re_compile_maybe_verbose(regex: str) -> Pattern[str]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3830
    """Compile a regular expression string in `regex`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3831
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3832
    If it contains newlines, use verbose mode.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3833
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3834
    if "\n" in regex:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3835
        regex = "(?x)" + regex
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3836
    return re.compile(regex)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3837
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3838
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3839
def enumerate_reversed(sequence: Sequence[T]) -> Iterator[Tuple[Index, T]]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3840
    """Like `reversed(enumerate(sequence))` if that were possible."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3841
    index = len(sequence) - 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3842
    for element in reversed(sequence):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3843
        yield (index, element)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3844
        index -= 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3845
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3846
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3847
def enumerate_with_length(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3848
    line: Line, reversed: bool = False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3849
) -> Iterator[Tuple[Index, Leaf, int]]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3850
    """Return an enumeration of leaves with their length.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3851
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3852
    Stops prematurely on multiline strings and standalone comments.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3853
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3854
    op = cast(
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3855
        Callable[[Sequence[Leaf]], Iterator[Tuple[Index, Leaf]]],
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3856
        enumerate_reversed if reversed else enumerate,
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3857
    )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3858
    for index, leaf in op(line.leaves):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3859
        length = len(leaf.prefix) + len(leaf.value)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3860
        if "\n" in leaf.value:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3861
            return  # Multiline strings, we can't continue.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3862
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3863
        for comment in line.comments_after(leaf):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3864
            length += len(comment.value)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3865
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3866
        yield index, leaf, length
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3867
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3868
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3869
def is_line_short_enough(line: Line, *, line_length: int, line_str: str = "") -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3870
    """Return True if `line` is no longer than `line_length`.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3871
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3872
    Uses the provided `line_str` rendering, if any, otherwise computes a new one.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3873
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3874
    if not line_str:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3875
        line_str = str(line).strip("\n")
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3876
    return (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3877
        len(line_str) <= line_length
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3878
        and "\n" not in line_str  # multiline strings
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3879
        and not line.contains_standalone_comments()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3880
    )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3881
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3882
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3883
def can_be_split(line: Line) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3884
    """Return False if the line cannot be split *for sure*.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3885
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3886
    This is not an exhaustive search but a cheap heuristic that we can use to
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3887
    avoid some unfortunate formattings (mostly around wrapping unsplittable code
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3888
    in unnecessary parentheses).
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3889
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3890
    leaves = line.leaves
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3891
    if len(leaves) < 2:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3892
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3893
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3894
    if leaves[0].type == token.STRING and leaves[1].type == token.DOT:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3895
        call_count = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3896
        dot_count = 0
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3897
        next = leaves[-1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3898
        for leaf in leaves[-2::-1]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3899
            if leaf.type in OPENING_BRACKETS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3900
                if next.type not in CLOSING_BRACKETS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3901
                    return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3902
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3903
                call_count += 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3904
            elif leaf.type == token.DOT:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3905
                dot_count += 1
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3906
            elif leaf.type == token.NAME:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3907
                if not (next.type == token.DOT or next.type in OPENING_BRACKETS):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3908
                    return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3909
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3910
            elif leaf.type not in CLOSING_BRACKETS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3911
                return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3912
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3913
            if dot_count > 1 and call_count > 1:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3914
                return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3915
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3916
    return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3917
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3918
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3919
def can_omit_invisible_parens(line: Line, line_length: int) -> bool:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3920
    """Does `line` have a shape safe to reformat without optional parens around it?
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3921
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3922
    Returns True for only a subset of potentially nice looking formattings but
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3923
    the point is to not return false positives that end up producing lines that
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3924
    are too long.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3925
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3926
    bt = line.bracket_tracker
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3927
    if not bt.delimiters:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3928
        # Without delimiters the optional parentheses are useless.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3929
        return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3930
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3931
    max_priority = bt.max_delimiter_priority()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3932
    if bt.delimiter_count_with_priority(max_priority) > 1:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3933
        # With more than one delimiter of a kind the optional parentheses read better.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3934
        return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3935
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3936
    if max_priority == DOT_PRIORITY:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3937
        # A single stranded method call doesn't require optional parentheses.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3938
        return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3939
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3940
    assert len(line.leaves) >= 2, "Stranded delimiter"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3941
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3942
    first = line.leaves[0]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3943
    second = line.leaves[1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3944
    penultimate = line.leaves[-2]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3945
    last = line.leaves[-1]
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3946
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3947
    # With a single delimiter, omit if the expression starts or ends with
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3948
    # a bracket.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3949
    if first.type in OPENING_BRACKETS and second.type not in CLOSING_BRACKETS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3950
        remainder = False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3951
        length = 4 * line.depth
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3952
        for _index, leaf, leaf_length in enumerate_with_length(line):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3953
            if leaf.type in CLOSING_BRACKETS and leaf.opening_bracket is first:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3954
                remainder = True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3955
            if remainder:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3956
                length += leaf_length
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3957
                if length > line_length:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3958
                    break
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3959
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3960
                if leaf.type in OPENING_BRACKETS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3961
                    # There are brackets we can further split on.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3962
                    remainder = False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3963
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3964
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3965
            # checked the entire string and line length wasn't exceeded
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3966
            if len(line.leaves) == _index + 1:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3967
                return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3968
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3969
        # Note: we are not returning False here because a line might have *both*
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3970
        # a leading opening bracket and a trailing closing bracket.  If the
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3971
        # opening bracket doesn't match our rule, maybe the closing will.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3972
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3973
    if (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3974
        last.type == token.RPAR
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3975
        or last.type == token.RBRACE
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3976
        or (
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3977
            # don't use indexing for omitting optional parentheses;
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3978
            # it looks weird
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3979
            last.type == token.RSQB
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3980
            and last.parent
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3981
            and last.parent.type != syms.trailer
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3982
        )
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3983
    ):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3984
        if penultimate.type in OPENING_BRACKETS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3985
            # Empty brackets don't help.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3986
            return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3987
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3988
        if is_multiline_string(first):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3989
            # Additional wrapping of a multiline string in this situation is
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3990
            # unnecessary.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3991
            return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3992
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3993
        length = 4 * line.depth
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3994
        seen_other_brackets = False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3995
        for _index, leaf, leaf_length in enumerate_with_length(line):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3996
            length += leaf_length
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3997
            if leaf is last.opening_bracket:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3998
                if seen_other_brackets or length <= line_length:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  3999
                    return True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4000
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4001
            elif leaf.type in OPENING_BRACKETS:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4002
                # There are brackets we can further split on.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4003
                seen_other_brackets = True
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4004
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4005
    return False
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4006
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4007
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4008
def get_cache_file(mode: FileMode) -> Path:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4009
    return CACHE_DIR / f"cache.{mode.get_cache_key()}.pickle"
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4010
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4011
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4012
def read_cache(mode: FileMode) -> Cache:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4013
    """Read the cache if it exists and is well formed.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4014
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4015
    If it is not well formed, the call to write_cache later should resolve the issue.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4016
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4017
    cache_file = get_cache_file(mode)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4018
    if not cache_file.exists():
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4019
        return {}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4020
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4021
    with cache_file.open("rb") as fobj:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4022
        try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4023
            cache: Cache = pickle.load(fobj)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4024
        except pickle.UnpicklingError:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4025
            return {}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4026
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4027
    return cache
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4028
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4029
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4030
def get_cache_info(path: Path) -> CacheInfo:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4031
    """Return the information used to check if a file is already formatted or not."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4032
    stat = path.stat()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4033
    return stat.st_mtime, stat.st_size
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4034
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4035
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4036
def filter_cached(cache: Cache, sources: Iterable[Path]) -> Tuple[Set[Path], Set[Path]]:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4037
    """Split an iterable of paths in `sources` into two sets.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4038
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4039
    The first contains paths of files that modified on disk or are not in the
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4040
    cache. The other contains paths to non-modified files.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4041
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4042
    todo, done = set(), set()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4043
    for src in sources:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4044
        src = src.resolve()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4045
        if cache.get(src) != get_cache_info(src):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4046
            todo.add(src)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4047
        else:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4048
            done.add(src)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4049
    return todo, done
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4050
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4051
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4052
def write_cache(cache: Cache, sources: Iterable[Path], mode: FileMode) -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4053
    """Update the cache file."""
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4054
    cache_file = get_cache_file(mode)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4055
    try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4056
        CACHE_DIR.mkdir(parents=True, exist_ok=True)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4057
        new_cache = {**cache, **{src.resolve(): get_cache_info(src) for src in sources}}
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4058
        with tempfile.NamedTemporaryFile(dir=str(cache_file.parent), delete=False) as f:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4059
            pickle.dump(new_cache, f, protocol=pickle.HIGHEST_PROTOCOL)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4060
        os.replace(f.name, cache_file)
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4061
    except OSError:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4062
        pass
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4063
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4064
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4065
def patch_click() -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4066
    """Make Click not crash.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4067
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4068
    On certain misconfigured environments, Python 3 selects the ASCII encoding as the
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4069
    default which restricts paths that it can access during the lifetime of the
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4070
    application.  Click refuses to work in this scenario by raising a RuntimeError.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4071
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4072
    In case of Black the likelihood that non-ASCII characters are going to be used in
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4073
    file paths is minimal since it's Python source code.  Moreover, this crash was
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4074
    spurious on Python 3.7 thanks to PEP 538 and PEP 540.
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4075
    """
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4076
    try:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4077
        from click import core
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4078
        from click import _unicodefun  # type: ignore
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4079
    except ModuleNotFoundError:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4080
        return
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4081
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4082
    for module in (core, _unicodefun):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4083
        if hasattr(module, "_verify_python3_env"):
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4084
            module._verify_python3_env = lambda: None
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4085
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4086
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4087
def patched_main() -> None:
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4088
    freeze_support()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4089
    patch_click()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4090
    main()
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4091
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4092
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4093
if __name__ == "__main__":
7054fd370430 contrib: add a fork of black (as "grey") that includes my changes
Augie Fackler <augie@google.com>
parents:
diff changeset
  4094
    patched_main()