| | from dataclasses import dataclass |
| | from typing import Iterator, TypeVar, Union |
| |
|
| | from black.nodes import Visitor |
| | from black.output import out |
| | from black.parsing import lib2to3_parse |
| | from blib2to3.pgen2 import token |
| | from blib2to3.pytree import Leaf, Node, type_repr |
| |
|
| | LN = Union[Leaf, Node] |
| | T = TypeVar("T") |
| |
|
| |
|
| | @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: |
| | |
| | |
| | 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)) |
| |
|
