1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54
---- ---- ---- ---- ---- ---- ---- -- -- --- --- --- --- --- --- ------- ------- ------- | from dataclasses import dataclass, field ---- ---- ---- ---- ---- ---- ---- -- -- --- --- --- --- --- --- ------- ------- ------- | from typing import Any, Iterator, List, 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 ---- ---- ---- ---- ---- ---- ---- 04 -- --- --- --- --- --- --- ------- ------- ------- | class DebugVisitor(Visitor[T]): ---- ---- ---- ---- ---- ---- ---- 04 -- --- --- --- --- --- --- ------- ------- ------- | tree_depth: int = 0 ---- ---- ---- ---- ---- ---- ---- 04 -- --- --- --- --- --- --- ------- ------- ------- | list_output: List[str] = field(default_factory=list) ---- ---- ---- ---- ---- ---- ---- 04 -- --- --- --- --- --- --- ------- ------- ------- | print_output: bool = True ---- ---- ---- ---- ---- ---- ---- 04 -- --- --- --- --- --- --- ------- ------- ------- | 0004 0004 0004 0000 0000 0000 0000 04 02 000 000 000 000 000 000 0000.00 0000.00 0000.00 | def out(self, message: str, *args: Any, **kwargs: Any) -> None: 0004 0004 0004 0000 0000 0000 0000 04 02 000 000 000 000 000 000 0000.00 0000.00 0000.00 | self.list_output.append(message) 0004 0004 0004 0000 0000 0000 0000 04 02 000 000 000 000 000 000 0000.00 0000.00 0000.00 | if self.print_output: 0004 0004 0004 0000 0000 0000 0000 04 02 000 000 000 000 000 000 0000.00 0000.00 0000.00 | out(message, *args, **kwargs) ---- ---- ---- ---- ---- ---- ---- 04 -- --- --- --- --- --- --- ------- ------- ------- | 0019 0016 0016 0002 0000 0001 0002 04 04 003 005 004 008 008 012 0036.00 0086.40 0002.40 | def visit_default(self, node: LN) -> Iterator[T]: 0019 0016 0016 0002 0000 0001 0002 04 04 003 005 004 008 008 012 0036.00 0086.40 0002.40 | indent = " " * (2 * self.tree_depth) 0019 0016 0016 0002 0000 0001 0002 04 04 003 005 004 008 008 012 0036.00 0086.40 0002.40 | if isinstance(node, Node): 0019 0016 0016 0002 0000 0001 0002 04 04 003 005 004 008 008 012 0036.00 0086.40 0002.40 | _type = type_repr(node.type) 0019 0016 0016 0002 0000 0001 0002 04 04 003 005 004 008 008 012 0036.00 0086.40 0002.40 | self.out(f"{indent}{_type}", fg="yellow") 0019 0016 0016 0002 0000 0001 0002 04 04 003 005 004 008 008 012 0036.00 0086.40 0002.40 | self.tree_depth += 1 0019 0016 0016 0002 0000 0001 0002 04 04 003 005 004 008 008 012 0036.00 0086.40 0002.40 | for child in node.children: 0019 0016 0016 0002 0000 0001 0002 04 04 003 005 004 008 008 012 0036.00 0086.40 0002.40 | yield from self.visit(child) 0019 0016 0016 0002 0000 0001 0002 04 04 003 005 004 008 008 012 0036.00 0086.40 0002.40 | 0019 0016 0016 0002 0000 0001 0002 04 04 003 005 004 008 008 012 0036.00 0086.40 0002.40 | self.tree_depth -= 1 0019 0016 0016 0002 0000 0001 0002 04 04 003 005 004 008 008 012 0036.00 0086.40 0002.40 | self.out(f"{indent}/{_type}", fg="yellow", bold=False) 0019 0016 0016 0002 0000 0001 0002 04 04 003 005 004 008 008 012 0036.00 0086.40 0002.40 | else: 0019 0016 0016 0002 0000 0001 0002 04 04 003 005 004 008 008 012 0036.00 0086.40 0002.40 | _type = token.tok_name.get(node.type, str(node.type)) 0019 0016 0016 0002 0000 0001 0002 04 04 003 005 004 008 008 012 0036.00 0086.40 0002.40 | self.out(f"{indent}{_type}", fg="blue", nl=False) 0019 0016 0016 0002 0000 0001 0002 04 04 003 005 004 008 008 012 0036.00 0086.40 0002.40 | if node.prefix: 0019 0016 0016 0002 0000 0001 0002 04 04 003 005 004 008 008 012 0036.00 0086.40 0002.40 | # We don't have to handle prefixes for `Node` objects since 0019 0016 0016 0002 0000 0001 0002 04 04 003 005 004 008 008 012 0036.00 0086.40 0002.40 | # that delegates to the first child anyway. 0019 0016 0016 0002 0000 0001 0002 04 04 003 005 004 008 008 012 0036.00 0086.40 0002.40 | self.out(f" {node.prefix!r}", fg="green", bold=False, nl=False) 0019 0016 0016 0002 0000 0001 0002 04 04 003 005 004 008 008 012 0036.00 0086.40 0002.40 | self.out(f" {node.value!r}", fg="blue", bold=False) ---- ---- ---- ---- ---- ---- ---- 04 -- --- --- --- --- --- --- ------- ------- ------- | ---- ---- ---- ---- ---- ---- ---- 04 -- --- --- --- --- --- --- ------- ------- ------- | @classmethod 0010 0008 0006 0000 0003 0001 0000 04 02 000 000 000 000 000 000 0000.00 0000.00 0000.00 | def show(cls, code: Union[str, Leaf, Node]) -> None: 0010 0008 0006 0000 0003 0001 0000 04 02 000 000 000 000 000 000 0000.00 0000.00 0000.00 | """Pretty-print the lib2to3 AST of a given string of `code`. 0010 0008 0006 0000 0003 0001 0000 04 02 000 000 000 000 000 000 0000.00 0000.00 0000.00 | 0010 0008 0006 0000 0003 0001 0000 04 02 000 000 000 000 000 000 0000.00 0000.00 0000.00 | Convenience method for debugging. 0010 0008 0006 0000 0003 0001 0000 04 02 000 000 000 000 000 000 0000.00 0000.00 0000.00 | """ 0010 0008 0006 0000 0003 0001 0000 04 02 000 000 000 000 000 000 0000.00 0000.00 0000.00 | v: DebugVisitor[None] = DebugVisitor() 0010 0008 0006 0000 0003 0001 0000 04 02 000 000 000 000 000 000 0000.00 0000.00 0000.00 | if isinstance(code, str): 0010 0008 0006 0000 0003 0001 0000 04 02 000 000 000 000 000 000 0000.00 0000.00 0000.00 | code = lib2to3_parse(code) 0010 0008 0006 0000 0003 0001 0000 04 02 000 000 000 000 000 000 0000.00 0000.00 0000.00 | list(v.visit(code))