import ast

import asyncio

from concurrent.futures import Executor, ProcessPoolExecutor

from contextlib import contextmanager

from datetime import datetime

from enum import Enum

from functools import lru_cache, partial, wraps

import io

import itertools

import logging

from multiprocessing import Manager, freeze_support

import os

from pathlib import Path

import pickle

import regex as re

import signal

import sys

import tempfile

import tokenize

import traceback

from typing import (

Any,

Callable,

Collection,

Dict,

Generator,

Generic,

Iterable,

Iterator,

List,

Optional,

Pattern,

Sequence,

Set,

Tuple,

TypeVar,

Union,

cast,

)

from typing_extensions import Final

from mypy_extensions import mypyc_attr

from appdirs import user_cache_dir

from dataclasses import dataclass, field, replace

import click

import toml

from typed_ast import ast3, ast27

from pathspec import PathSpec

# lib2to3 fork

from blib2to3.pytree import Node, Leaf, type_repr

from blib2to3 import pygram, pytree

from blib2to3.pgen2 import driver, token

from blib2to3.pgen2.grammar import Grammar

from blib2to3.pgen2.parse import ParseError

from _black_version import version as __version__

DEFAULT_LINE_LENGTH = 88

DEFAULT_EXCLUDES = r"/(\.eggs|\.git|\.hg|\.mypy_cache|\.nox|\.tox|\.venv|\.svn|_build|buck-out|build|dist)/" # noqa: B950

DEFAULT_INCLUDES = r"\.pyi?$"

CACHE_DIR = Path(user_cache_dir("black", version=__version__))

# types

FileContent = str

Encoding = str

NewLine = str

Depth = int

NodeType = int

LeafID = int

Priority = int

Index = int

LN = Union[Leaf, Node]

SplitFunc = Callable[["Line", Collection["Feature"]], Iterator["Line"]]

Timestamp = float

FileSize = int

CacheInfo = Tuple[Timestamp, FileSize]

Cache = Dict[Path, CacheInfo]

out = partial(click.secho, bold=True, err=True)

err = partial(click.secho, fg="red", err=True)

pygram.initialize(CACHE_DIR)

syms = pygram.python_symbols

class NothingChanged(UserWarning):

"""Raised when reformatted code is the same as source."""

class CannotSplit(Exception):

"""A readable split that fits the allotted line length is impossible."""

class InvalidInput(ValueError):

"""Raised when input source code fails all parse attempts."""

class WriteBack(Enum):

NO = 0

YES = 1

DIFF = 2

CHECK = 3

@classmethod

def from_configuration(cls, *, check: bool, diff: bool) -> "WriteBack":

if check and not diff:

return cls.CHECK

return cls.DIFF if diff else cls.YES

class Changed(Enum):

NO = 0

CACHED = 1

YES = 2

class TargetVersion(Enum):

PY27 = 2

PY33 = 3

PY34 = 4

PY35 = 5

PY36 = 6

PY37 = 7

PY38 = 8

def is_python2(self) -> bool:

return self is TargetVersion.PY27

PY36_VERSIONS = {TargetVersion.PY36, TargetVersion.PY37, TargetVersion.PY38}

class Feature(Enum):

# All string literals are unicode

UNICODE_LITERALS = 1

F_STRINGS = 2

NUMERIC_UNDERSCORES = 3

TRAILING_COMMA_IN_CALL = 4

TRAILING_COMMA_IN_DEF = 5

# The following two feature-flags are mutually exclusive, and exactly one should be

# set for every version of python.

ASYNC_IDENTIFIERS = 6

ASYNC_KEYWORDS = 7

ASSIGNMENT_EXPRESSIONS = 8

POS_ONLY_ARGUMENTS = 9

VERSION_TO_FEATURES: Dict[TargetVersion, Set[Feature]] = {

TargetVersion.PY27: {Feature.ASYNC_IDENTIFIERS},

TargetVersion.PY33: {Feature.UNICODE_LITERALS, Feature.ASYNC_IDENTIFIERS},

TargetVersion.PY34: {Feature.UNICODE_LITERALS, Feature.ASYNC_IDENTIFIERS},

TargetVersion.PY35: {

Feature.UNICODE_LITERALS,

Feature.TRAILING_COMMA_IN_CALL,

Feature.ASYNC_IDENTIFIERS,

},

TargetVersion.PY36: {

Feature.UNICODE_LITERALS,

Feature.F_STRINGS,

Feature.NUMERIC_UNDERSCORES,

Feature.TRAILING_COMMA_IN_CALL,

Feature.TRAILING_COMMA_IN_DEF,

Feature.ASYNC_IDENTIFIERS,

},

TargetVersion.PY37: {

Feature.UNICODE_LITERALS,

Feature.F_STRINGS,

Feature.NUMERIC_UNDERSCORES,

Feature.TRAILING_COMMA_IN_CALL,

Feature.TRAILING_COMMA_IN_DEF,

Feature.ASYNC_KEYWORDS,

},

TargetVersion.PY38: {

Feature.UNICODE_LITERALS,

Feature.F_STRINGS,

Feature.NUMERIC_UNDERSCORES,

Feature.TRAILING_COMMA_IN_CALL,

Feature.TRAILING_COMMA_IN_DEF,

Feature.ASYNC_KEYWORDS,

Feature.ASSIGNMENT_EXPRESSIONS,

Feature.POS_ONLY_ARGUMENTS,

},

}

@dataclass

class FileMode:

target_versions: Set[TargetVersion] = field(default_factory=set)

line_length: int = DEFAULT_LINE_LENGTH

string_normalization: bool = True

is_pyi: bool = False

def get_cache_key(self) -> str:

if self.target_versions:

version_str = ",".join(

str(version.value)

for version in sorted(self.target_versions, key=lambda v: v.value)

)

else:

version_str = "-"

parts = [

version_str,

str(self.line_length),

str(int(self.string_normalization)),

str(int(self.is_pyi)),

]

return ".".join(parts)

def supports_feature(target_versions: Set[TargetVersion], feature: Feature) -> bool:

return all(feature in VERSION_TO_FEATURES[version] for version in target_versions)

def find_pyproject_toml(path_search_start: str) -> Optional[str]:

"""Find the absolute filepath to a pyproject.toml if it exists"""

path_project_root = find_project_root(path_search_start)

path_pyproject_toml = path_project_root / "pyproject.toml"

return str(path_pyproject_toml) if path_pyproject_toml.is_file() else None

def parse_pyproject_toml(path_config: str) -> Dict[str, Any]:

"""Parse a pyproject toml file, pulling out relevant parts for Black

If parsing fails, will raise a toml.TomlDecodeError

"""

pyproject_toml = toml.load(path_config)

config = pyproject_toml.get("tool", {}).get("black", {})

return {k.replace("--", "").replace("-", "_"): v for k, v in config.items()}

def read_pyproject_toml(

ctx: click.Context, param: click.Parameter, value: Union[str, int, bool, None]

) -> Optional[str]:

"""Inject Black configuration from "pyproject.toml" into defaults in `ctx`.

Returns the path to a successfully found and read configuration file, None

otherwise.

"""

assert not isinstance(value, (int, bool)), "Invalid parameter type passed"

if not value:

value = find_pyproject_toml(ctx.params.get("src", ()))

if value is None:

return None

try:

config = parse_pyproject_toml(value)

except (toml.TomlDecodeError, OSError) as e:

raise click.FileError(

filename=value, hint=f"Error reading configuration file: {e}"

)

if not config:

return None

if ctx.default_map is None:

ctx.default_map = {}

ctx.default_map.update(config) # type: ignore # bad types in .pyi

return value

def target_version_option_callback(

c: click.Context, p: Union[click.Option, click.Parameter], v: Tuple[str, ...]

) -> List[TargetVersion]:

"""Compute the target versions from a --target-version flag.

This is its own function because mypy couldn't infer the type correctly

when it was a lambda, causing mypyc trouble.

"""

return [TargetVersion[val.upper()] for val in v]

@click.command(context_settings=dict(help_option_names=["-h", "--help"]))

@click.option("-c", "--code", type=str, help="Format the code passed in as a string.")

@click.option(

"-l",

"--line-length",

type=int,

default=DEFAULT_LINE_LENGTH,

help="How many characters per line to allow.",

show_default=True,

)

@click.option(

"-t",

"--target-version",

type=click.Choice([v.name.lower() for v in TargetVersion]),

callback=target_version_option_callback,

multiple=True,

help=(

"Python versions that should be supported by Black's output. [default: "

"per-file auto-detection]"

),

)

@click.option(

"--py36",

is_flag=True,

help=(

"Allow using Python 3.6-only syntax on all input files. This will put "

"trailing commas in function signatures and calls also after *args and "

"**kwargs. Deprecated; use --target-version instead. "

"[default: per-file auto-detection]"

),

)

@click.option(

"--pyi",

is_flag=True,

help=(

"Format all input files like typing stubs regardless of file extension "

"(useful when piping source on standard input)."

),

)

@click.option(

"-S",

"--skip-string-normalization",

is_flag=True,

help="Don't normalize string quotes or prefixes.",

)

@click.option(

"--check",

is_flag=True,

help=(

"Don't write the files back, just return the status. Return code 0 "

"means nothing would change. Return code 1 means some files would be "

"reformatted. Return code 123 means there was an internal error."

),

)

@click.option(

"--diff",

is_flag=True,

help="Don't write the files back, just output a diff for each file on stdout.",

)

@click.option(

"--fast/--safe",

is_flag=True,

help="If --fast given, skip temporary sanity checks. [default: --safe]",

)

@click.option(

"--include",

type=str,

default=DEFAULT_INCLUDES,

help=(

"A regular expression that matches files and directories that should be "

"included on recursive searches. An empty value means all files are "

"included regardless of the name. Use forward slashes for directories on "

"all platforms (Windows, too). Exclusions are calculated first, inclusions "

"later."

),

show_default=True,

)

@click.option(

"--exclude",

type=str,

default=DEFAULT_EXCLUDES,

help=(

"A regular expression that matches files and directories that should be "

"excluded on recursive searches. An empty value means no paths are excluded. "

"Use forward slashes for directories on all platforms (Windows, too). "

"Exclusions are calculated first, inclusions later."

),

show_default=True,

)

@click.option(

"-q",

"--quiet",

is_flag=True,

help=(

"Don't emit non-error messages to stderr. Errors are still emitted; "

"silence those with 2>/dev/null."

),

)

@click.option(

"-v",

"--verbose",

is_flag=True,

help=(

"Also emit messages to stderr about files that were not changed or were "

"ignored due to --exclude=."

),

)

@click.version_option(version=__version__)

@click.argument(

"src",

nargs=-1,

type=click.Path(

exists=True, file_okay=True, dir_okay=True, readable=True, allow_dash=True

),

is_eager=True,

)

@click.option(

"--config",

type=click.Path(

exists=False, file_okay=True, dir_okay=False, readable=True, allow_dash=False

),

is_eager=True,

callback=read_pyproject_toml,

help="Read configuration from PATH.",

)

@click.pass_context

def main(

ctx: click.Context,

code: Optional[str],

line_length: int,

target_version: List[TargetVersion],

check: bool,

diff: bool,

fast: bool,

pyi: bool,

py36: bool,

skip_string_normalization: bool,

quiet: bool,

verbose: bool,

include: str,

exclude: str,

src: Tuple[str, ...],

config: Optional[str],

) -> None:

"""The uncompromising code formatter."""

write_back = WriteBack.from_configuration(check=check, diff=diff)

if target_version:

if py36:

err(f"Cannot use both --target-version and --py36")

ctx.exit(2)

else:

versions = set(target_version)

elif py36:

err(

"--py36 is deprecated and will be removed in a future version. "

"Use --target-version py36 instead."

)

versions = PY36_VERSIONS

else:

# We'll autodetect later.

versions = set()

mode = FileMode(

target_versions=versions,

line_length=line_length,

is_pyi=pyi,

string_normalization=not skip_string_normalization,

)

if config and verbose:

out(f"Using configuration from {config}.", bold=False, fg="blue")

if code is not None:

print(format_str(code, mode=mode))

ctx.exit(0)

try:

include_regex = re_compile_maybe_verbose(include)

except re.error:

err(f"Invalid regular expression for include given: {include!r}")

ctx.exit(2)

try:

exclude_regex = re_compile_maybe_verbose(exclude)

except re.error:

err(f"Invalid regular expression for exclude given: {exclude!r}")

ctx.exit(2)

report = Report(check=check, diff=diff, quiet=quiet, verbose=verbose)

root = find_project_root(src)

sources: Set[Path] = set()

path_empty(src, quiet, verbose, ctx)

for s in src:

p = Path(s)

if p.is_dir():

sources.update(

gen_python_files_in_dir(

p, root, include_regex, exclude_regex, report, get_gitignore(root)

)

)

elif p.is_file() or s == "-":

# if a file was explicitly given, we don't care about its extension

sources.add(p)

else:

err(f"invalid path: {s}")

if len(sources) == 0:

if verbose or not quiet:

out("No Python files are present to be formatted. Nothing to do 😴")

ctx.exit(0)

if len(sources) == 1:

reformat_one(

src=sources.pop(),

fast=fast,

write_back=write_back,

mode=mode,

report=report,

)

else:

reformat_many(

sources=sources, fast=fast, write_back=write_back, mode=mode, report=report

)

if verbose or not quiet:

out("Oh no! 💥 💔 💥" if report.return_code else "All done! ✨ 🍰 ✨")

click.secho(str(report), err=True)

ctx.exit(report.return_code)

def path_empty(

src: Tuple[str, ...], quiet: bool, verbose: bool, ctx: click.Context

) -> None:

"""

Exit if there is no `src` provided for formatting

"""

if not src:

if verbose or not quiet:

out("No Path provided. Nothing to do 😴")

ctx.exit(0)

def reformat_one(

src: Path, fast: bool, write_back: WriteBack, mode: FileMode, report: "Report"

) -> None:

"""Reformat a single file under `src` without spawning child processes.

`fast`, `write_back`, and `mode` options are passed to

:func:`format_file_in_place` or :func:`format_stdin_to_stdout`.

"""

try:

changed = Changed.NO

if not src.is_file() and str(src) == "-":

if format_stdin_to_stdout(fast=fast, write_back=write_back, mode=mode):

changed = Changed.YES

else:

cache: Cache = {}

if write_back != WriteBack.DIFF:

cache = read_cache(mode)

res_src = src.resolve()

if res_src in cache and cache[res_src] == get_cache_info(res_src):

changed = Changed.CACHED

if changed is not Changed.CACHED and format_file_in_place(

src, fast=fast, write_back=write_back, mode=mode

):

changed = Changed.YES

if (write_back is WriteBack.YES and changed is not Changed.CACHED) or (

write_back is WriteBack.CHECK and changed is Changed.NO

):

write_cache(cache, [src], mode)

report.done(src, changed)

except Exception as exc:

report.failed(src, str(exc))

def reformat_many(

sources: Set[Path],

fast: bool,

write_back: WriteBack,

mode: FileMode,

report: "Report",

) -> None:

"""Reformat multiple files using a ProcessPoolExecutor."""

loop = asyncio.get_event_loop()

worker_count = os.cpu_count()

if sys.platform == "win32":

# Work around https://bugs.python.org/issue26903

worker_count = min(worker_count, 61)

executor = ProcessPoolExecutor(max_workers=worker_count)

try:

loop.run_until_complete(

schedule_formatting(

sources=sources,

fast=fast,

write_back=write_back,

mode=mode,

report=report,

loop=loop,

executor=executor,

)

)

finally:

shutdown(loop)

executor.shutdown()

async def schedule_formatting(

sources: Set[Path],

fast: bool,

write_back: WriteBack,

mode: FileMode,

report: "Report",

loop: asyncio.AbstractEventLoop,

executor: Executor,

) -> None:

"""Run formatting of `sources` in parallel using the provided `executor`.

(Use ProcessPoolExecutors for actual parallelism.)

`write_back`, `fast`, and `mode` options are passed to

:func:`format_file_in_place`.

"""

cache: Cache = {}

if write_back != WriteBack.DIFF:

cache = read_cache(mode)

sources, cached = filter_cached(cache, sources)

for src in sorted(cached):

report.done(src, Changed.CACHED)

if not sources:

return

cancelled = []

sources_to_cache = []

lock = None

if write_back == WriteBack.DIFF:

# For diff output, we need locks to ensure we don't interleave output

# from different processes.

manager = Manager()

lock = manager.Lock()

tasks = {

asyncio.ensure_future(

loop.run_in_executor(

executor, format_file_in_place, src, fast, mode, write_back, lock

)

): src

for src in sorted(sources)

}

pending: Iterable["asyncio.Future[bool]"] = tasks.keys()

try:

loop.add_signal_handler(signal.SIGINT, cancel, pending)

loop.add_signal_handler(signal.SIGTERM, cancel, pending)

except NotImplementedError:

# There are no good alternatives for these on Windows.

pass

while pending:

done, _ = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)

for task in done:

src = tasks.pop(task)

if task.cancelled():

cancelled.append(task)

elif task.exception():

report.failed(src, str(task.exception()))

else:

changed = Changed.YES if task.result() else Changed.NO

# If the file was written back or was successfully checked as

# well-formatted, store this information in the cache.

if write_back is WriteBack.YES or (

write_back is WriteBack.CHECK and changed is Changed.NO

):

sources_to_cache.append(src)

report.done(src, changed)

if cancelled:

await asyncio.gather(*cancelled, loop=loop, return_exceptions=True)

if sources_to_cache:

write_cache(cache, sources_to_cache, mode)

def format_file_in_place(

src: Path,

fast: bool,

mode: FileMode,

write_back: WriteBack = WriteBack.NO,

lock: Any = None, # multiprocessing.Manager().Lock() is some crazy proxy

) -> bool:

"""Format file under `src` path. Return True if changed.

If `write_back` is DIFF, write a diff to stdout. If it is YES, write reformatted

code to the file.

`mode` and `fast` options are passed to :func:`format_file_contents`.

"""

if src.suffix == ".pyi":

mode = replace(mode, is_pyi=True)

then = datetime.utcfromtimestamp(src.stat().st_mtime)

with open(src, "rb") as buf:

src_contents, encoding, newline = decode_bytes(buf.read())

try:

dst_contents = format_file_contents(src_contents, fast=fast, mode=mode)

except NothingChanged:

return False

if write_back == WriteBack.YES:

with open(src, "w", encoding=encoding, newline=newline) as f:

f.write(dst_contents)

elif write_back == WriteBack.DIFF:

now = datetime.utcnow()

src_name = f"{src}\t{then} +0000"

dst_name = f"{src}\t{now} +0000"

diff_contents = diff(src_contents, dst_contents, src_name, dst_name)

with lock or nullcontext():

f = io.TextIOWrapper(

sys.stdout.buffer,

encoding=encoding,

newline=newline,

write_through=True,

)

f.write(diff_contents)

f.detach()

return True

def format_stdin_to_stdout(

fast: bool, *, write_back: WriteBack = WriteBack.NO, mode: FileMode

) -> bool:

"""Format file on stdin. Return True if changed.

If `write_back` is YES, write reformatted code back to stdout. If it is DIFF,

write a diff to stdout. The `mode` argument is passed to

:func:`format_file_contents`.

"""

then = datetime.utcnow()

src, encoding, newline = decode_bytes(sys.stdin.buffer.read())

dst = src

try:

dst = format_file_contents(src, fast=fast, mode=mode)

return True

except NothingChanged:

return False

finally:

f = io.TextIOWrapper(

sys.stdout.buffer, encoding=encoding, newline=newline, write_through=True

)

if write_back == WriteBack.YES:

f.write(dst)

elif write_back == WriteBack.DIFF:

now = datetime.utcnow()

src_name = f"STDIN\t{then} +0000"

dst_name = f"STDOUT\t{now} +0000"

f.write(diff(src, dst, src_name, dst_name))

f.detach()

def format_file_contents(

src_contents: str, *, fast: bool, mode: FileMode

) -> FileContent:

"""Reformat contents a file and return new contents.

If `fast` is False, additionally confirm that the reformatted code is

valid by calling :func:`assert_equivalent` and :func:`assert_stable` on it.

`mode` is passed to :func:`format_str`.

"""

if src_contents.strip() == "":

raise NothingChanged

dst_contents = format_str(src_contents, mode=mode)

if src_contents == dst_contents:

raise NothingChanged

if not fast:

assert_equivalent(src_contents, dst_contents)

assert_stable(src_contents, dst_contents, mode=mode)

return dst_contents

def format_str(src_contents: str, *, mode: FileMode) -> FileContent:

"""Reformat a string and return new contents.

`mode` determines formatting options, such as how many characters per line are

allowed.

"""

src_node = lib2to3_parse(src_contents.lstrip(), mode.target_versions)

dst_contents = []

future_imports = get_future_imports(src_node)

if mode.target_versions:

versions = mode.target_versions

else:

versions = detect_target_versions(src_node)

normalize_fmt_off(src_node)

lines = LineGenerator(

remove_u_prefix="unicode_literals" in future_imports

or supports_feature(versions, Feature.UNICODE_LITERALS),

is_pyi=mode.is_pyi,

normalize_strings=mode.string_normalization,

)

elt = EmptyLineTracker(is_pyi=mode.is_pyi)

empty_line = Line()

after = 0

split_line_features = {

feature

for feature in {Feature.TRAILING_COMMA_IN_CALL, Feature.TRAILING_COMMA_IN_DEF}

if supports_feature(versions, feature)

}

for current_line in lines.visit(src_node):

dst_contents.append(str(empty_line) * after)

before, after = elt.maybe_empty_lines(current_line)

dst_contents.append(str(empty_line) * before)

for line in split_line(

current_line, line_length=mode.line_length, features=split_line_features

):

dst_contents.append(str(line))

return "".join(dst_contents)

def decode_bytes(src: bytes) -> Tuple[FileContent, Encoding, NewLine]:

"""Return a tuple of (decoded_contents, encoding, newline).

`newline` is either CRLF or LF but `decoded_contents` is decoded with

universal newlines (i.e. only contains LF).

"""

srcbuf = io.BytesIO(src)

encoding, lines = tokenize.detect_encoding(srcbuf.readline)

if not lines:

return "", encoding, "\n"

newline = "\r\n" if b"\r\n" == lines[0][-2:] else "\n"

srcbuf.seek(0)

with io.TextIOWrapper(srcbuf, encoding) as tiow:

return tiow.read(), encoding, newline

def get_grammars(target_versions: Set[TargetVersion]) -> List[Grammar]:

if not target_versions:

# No target_version specified, so try all grammars.

return [

# Python 3.7+

pygram.python_grammar_no_print_statement_no_exec_statement_async_keywords,

# Python 3.0-3.6

pygram.python_grammar_no_print_statement_no_exec_statement,

# Python 2.7 with future print_function import

pygram.python_grammar_no_print_statement,

# Python 2.7

pygram.python_grammar,

]

if all(version.is_python2() for version in target_versions):

# Python 2-only code, so try Python 2 grammars.

return [

# Python 2.7 with future print_function import

pygram.python_grammar_no_print_statement,

# Python 2.7

pygram.python_grammar,

]

# Python 3-compatible code, so only try Python 3 grammar.

grammars = []

# If we have to parse both, try to parse async as a keyword first

if not supports_feature(target_versions, Feature.ASYNC_IDENTIFIERS):

# Python 3.7+

grammars.append(

pygram.python_grammar_no_print_statement_no_exec_statement_async_keywords

)

if not supports_feature(target_versions, Feature.ASYNC_KEYWORDS):

# Python 3.0-3.6

grammars.append(pygram.python_grammar_no_print_statement_no_exec_statement)

# At least one of the above branches must have been taken, because every Python

# version has exactly one of the two 'ASYNC_*' flags

return grammars

def lib2to3_parse(src_txt: str, target_versions: Iterable[TargetVersion] = ()) -> Node:

"""Given a string with source, return the lib2to3 Node."""

if src_txt[-1:] != "\n":

src_txt += "\n"

for grammar in get_grammars(set(target_versions)):

drv = driver.Driver(grammar, pytree.convert)

try:

result = drv.parse_string(src_txt, True)

break

except ParseError as pe:

lineno, column = pe.context[1]

lines = src_txt.splitlines()

try:

faulty_line = lines[lineno - 1]

except IndexError:

faulty_line = "<line number missing in source>"

exc = InvalidInput(f"Cannot parse: {lineno}:{column}: {faulty_line}")

else:

raise exc from None

if isinstance(result, Leaf):

result = Node(syms.file_input, [result])

return result

def lib2to3_unparse(node: Node) -> str:

"""Given a lib2to3 node, return its string representation."""

code = str(node)

return code

T = TypeVar("T")

class Visitor(Generic[T]):

"""Basic lib2to3 visitor that yields things of type `T` on `visit()`."""

def visit(self, node: LN) -> Iterator[T]:

"""Main method to visit `node` and its children.

It tries to find a `visit_*()` method for the given `node.type`, like

`visit_simple_stmt` for Node objects or `visit_INDENT` for Leaf objects.

If no dedicated `visit_*()` method is found, chooses `visit_default()`

instead.

Then yields objects of type `T` from the selected visitor.

"""

if node.type < 256:

name = token.tok_name[node.type]

else:

name = str(type_repr(node.type))

# We explicitly branch on whether a visitor exists (instead of

# using self.visit_default as the default arg to getattr) in order

# to save needing to create a bound method object and so mypyc can

# generate a native call to visit_default.

visitf = getattr(self, f"visit_{name}", None)

if visitf:

yield from visitf(node)

else:

yield from self.visit_default(node)

def visit_default(self, node: LN) -> Iterator[T]:

"""Default `visit_*()` implementation. Recurses to children of `node`."""

if isinstance(node, Node):

for child in node.children:

yield from self.visit(child)

@dataclass

class DebugVisitor(Visitor[T]):

tree_depth: int = 0

def visit_default(self, node: LN) -> Iterator[T]:

indent = " " * (2 * self.tree_depth)

if isinstance(node, Node):

_type = type_repr(node.type)

out(f"{indent}{_type}", fg="yellow")

self.tree_depth += 1

for child in node.children:

yield from self.visit(child)

self.tree_depth -= 1

out(f"{indent}/{_type}", fg="yellow", bold=False)

else:

_type = token.tok_name.get(node.type, str(node.type))

out(f"{indent}{_type}", fg="blue", nl=False)

if node.prefix:

# We don't have to handle prefixes for `Node` objects since

# that delegates to the first child anyway.

out(f" {node.prefix!r}", fg="green", bold=False, nl=False)

out(f" {node.value!r}", fg="blue", bold=False)

@classmethod

def show(cls, code: Union[str, Leaf, Node]) -> None:

"""Pretty-print the lib2to3 AST of a given string of `code`.

Convenience method for debugging.

"""

v: DebugVisitor[None] = DebugVisitor()

if isinstance(code, str):

code = lib2to3_parse(code)

list(v.visit(code))

WHITESPACE: Final = {token.DEDENT, token.INDENT, token.NEWLINE}

STATEMENT: Final = {

syms.if_stmt,

syms.while_stmt,

syms.for_stmt,

syms.try_stmt,

syms.except_clause,

syms.with_stmt,

syms.funcdef,

syms.classdef,

}

STANDALONE_COMMENT: Final = 153

token.tok_name[STANDALONE_COMMENT] = "STANDALONE_COMMENT"

LOGIC_OPERATORS: Final = {"and", "or"}

COMPARATORS: Final = {

token.LESS,

token.GREATER,

token.EQEQUAL,

token.NOTEQUAL,

token.LESSEQUAL,

token.GREATEREQUAL,

}

MATH_OPERATORS: Final = {

token.VBAR,

token.CIRCUMFLEX,

token.AMPER,

token.LEFTSHIFT,

token.RIGHTSHIFT,

token.PLUS,

token.MINUS,

token.STAR,

token.SLASH,

token.DOUBLESLASH,

token.PERCENT,

token.AT,

token.TILDE,

token.DOUBLESTAR,

}

STARS: Final = {token.STAR, token.DOUBLESTAR}

VARARGS_SPECIALS: Final = STARS | {token.SLASH}

VARARGS_PARENTS: Final = {

syms.arglist,

syms.argument, # double star in arglist

syms.trailer, # single argument to call

syms.typedargslist,

syms.varargslist, # lambdas

}

UNPACKING_PARENTS: Final = {

syms.atom, # single element of a list or set literal

syms.dictsetmaker,

syms.listmaker,

syms.testlist_gexp,

syms.testlist_star_expr,

}

TEST_DESCENDANTS: Final = {