104 lines
2.8 KiB
Python
104 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
import io
|
|
from contextlib import contextmanager
|
|
from enum import Enum, auto
|
|
from typing import Callable, Generator, Generic, Optional, Protocol, Sequence, TypeVar
|
|
|
|
|
|
class _Level(Enum):
|
|
EMPTY = auto()
|
|
ACTIVE = auto()
|
|
LAST = auto()
|
|
|
|
|
|
class Expr(Protocol):
|
|
def accept(self, printer: AstPrinter) -> None: ...
|
|
|
|
|
|
T = TypeVar("T", bound=Expr)
|
|
|
|
|
|
class AstPrinter(Generic[T]):
|
|
LAST_CHILD = "└── "
|
|
CHILD = "├── "
|
|
VERTICAL = "│ "
|
|
EMPTY = " "
|
|
|
|
def __init__(self):
|
|
self._levels: list[_Level] = []
|
|
self._idx: Optional[int] = None
|
|
self._buf: io.StringIO = io.StringIO()
|
|
|
|
def print(self, expr: T):
|
|
self._buf = io.StringIO()
|
|
expr.accept(self)
|
|
return self._buf.getvalue()
|
|
|
|
@contextmanager
|
|
def _child_level(self, single: bool = False) -> Generator[None, None, None]:
|
|
self._levels.append(_Level.LAST if single else _Level.ACTIVE)
|
|
try:
|
|
yield
|
|
finally:
|
|
self._levels.pop()
|
|
|
|
def _mark_last(self):
|
|
if self._levels:
|
|
self._levels[-1] = _Level.LAST
|
|
|
|
def _write_line(self, text: str, *, last: bool = False):
|
|
if last:
|
|
self._mark_last()
|
|
indent: str = self._build_indent()
|
|
if self._idx is not None:
|
|
text = f"[{self._idx}] {text}"
|
|
self._idx = None
|
|
self._buf.write(indent + text + "\n")
|
|
|
|
def _build_indent(self) -> str:
|
|
parts: list[str] = []
|
|
for level in self._levels[:-1]:
|
|
parts.append(self.EMPTY if level == _Level.EMPTY else self.VERTICAL)
|
|
if self._levels:
|
|
if self._levels[-1] == _Level.LAST:
|
|
parts.append(self.LAST_CHILD)
|
|
self._levels[-1] = _Level.EMPTY
|
|
else:
|
|
parts.append(self.CHILD)
|
|
return "".join(parts)
|
|
|
|
def _write_optional_child(
|
|
self, label: str, child: Optional[T], *, last: bool = False
|
|
):
|
|
if last:
|
|
self._mark_last()
|
|
if child is None:
|
|
self._write_line(f"{label}: None")
|
|
else:
|
|
self._write_line(label)
|
|
with self._child_level(single=True):
|
|
child.accept(self)
|
|
|
|
def _write_sequence(
|
|
self,
|
|
label: str,
|
|
list_: Sequence[T],
|
|
*,
|
|
last: bool = False,
|
|
print_func: Optional[Callable[[T], None]] = None,
|
|
):
|
|
if last:
|
|
self._mark_last()
|
|
|
|
self._write_line(label)
|
|
with self._child_level():
|
|
for i, item in enumerate(list_):
|
|
self._idx = i
|
|
if i == len(list_) - 1:
|
|
self._mark_last()
|
|
if print_func is not None:
|
|
print_func(item)
|
|
else:
|
|
item.accept(self)
|