]>
Commit | Line | Data |
---|---|---|
1 | from dataclasses import dataclass | |
2 | from typing import Iterator, TypeVar, Union | |
3 | ||
4 | from black.nodes import Visitor | |
5 | from black.output import out | |
6 | from black.parsing import lib2to3_parse | |
7 | from blib2to3.pgen2 import token | |
8 | from blib2to3.pytree import Leaf, Node, type_repr | |
9 | ||
10 | LN = Union[Leaf, Node] | |
11 | T = TypeVar("T") | |
12 | ||
13 | ||
14 | @dataclass | |
15 | class DebugVisitor(Visitor[T]): | |
16 | tree_depth: int = 0 | |
17 | ||
18 | def visit_default(self, node: LN) -> Iterator[T]: | |
19 | indent = " " * (2 * self.tree_depth) | |
20 | if isinstance(node, Node): | |
21 | _type = type_repr(node.type) | |
22 | out(f"{indent}{_type}", fg="yellow") | |
23 | self.tree_depth += 1 | |
24 | for child in node.children: | |
25 | yield from self.visit(child) | |
26 | ||
27 | self.tree_depth -= 1 | |
28 | out(f"{indent}/{_type}", fg="yellow", bold=False) | |
29 | else: | |
30 | _type = token.tok_name.get(node.type, str(node.type)) | |
31 | out(f"{indent}{_type}", fg="blue", nl=False) | |
32 | if node.prefix: | |
33 | # We don't have to handle prefixes for `Node` objects since | |
34 | # that delegates to the first child anyway. | |
35 | out(f" {node.prefix!r}", fg="green", bold=False, nl=False) | |
36 | out(f" {node.value!r}", fg="blue", bold=False) | |
37 | ||
38 | @classmethod | |
39 | def show(cls, code: Union[str, Leaf, Node]) -> None: | |
40 | """Pretty-print the lib2to3 AST of a given string of `code`. | |
41 | ||
42 | Convenience method for debugging. | |
43 | """ | |
44 | v: DebugVisitor[None] = DebugVisitor() | |
45 | if isinstance(code, str): | |
46 | code = lib2to3_parse(code) | |
47 | list(v.visit(code)) |