refactor(parser): improve AST printer

refactored the messy AST printer impletation with Claude to use a context manager, an enum and extract common functions

Co-authored-by: Claude <noreply@anthropic.com>

core/ast/printer.py

@@ -1,84 +1,97 @@
-from typing import Optional
+from contextlib import contextmanager
+from enum import Enum, auto
+import io
+from typing import Generator, Optional
 
 from core.ast.annotations import Expr, TypeExpr, SchemaExpr, SchemaElementExpr
 
 
-class AnnotationAstPrinter(Expr.Visitor[str]):
+class _Level(Enum):
+    EMPTY = auto()
+    ACTIVE = auto()
+    LAST = auto()
+
+
+class AnnotationAstPrinter(Expr.Visitor[None]):
     LAST_CHILD = "└── "
     CHILD = "├── "
     VERTICAL = "│   "
     EMPTY = "    "
 
     def __init__(self):
-        self.level: int = 0
+        self._levels: list[_Level] = []
-        self.idx: Optional[int] = None
+        self._idx: Optional[int] = None
-        self.last: bool = False
+        self._buf: io.StringIO = io.StringIO()
-        self.levels: list[int] = []
 
     def print(self, expr: Expr):
-        return expr.accept(self)
+        self._buf = io.StringIO()
+        expr.accept(self)
+        return self._buf.getvalue()
 
-    def print_line(self, text: str) -> str:
+    @contextmanager
-        indent: str = ""
+    def _child_level(self, last: bool = False) -> Generator[None, None, None]:
-        for enabled in self.levels[:-1]:
+        self._levels.append(_Level.LAST if last else _Level.ACTIVE)
-            if enabled:
+        try:
-                indent += self.VERTICAL
+            yield
+        finally:
+            self._levels.pop()
+
+    def _mark_last(self):
+        if self._levels:
+            self._levels[-1] = _Level.LAST
+
+    def _write_line(self, text: str):
+        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:
-                indent += self.EMPTY
+                parts.append(self.CHILD)
+        return "".join(parts)
 
-        if len(self.levels) > 0:
+    def _write_optional_child(
-            if self.levels[-1] == 2:
+        self, label: str, child: Optional[Expr], *, last: bool = False
-                indent += self.LAST_CHILD
+    ):
-                self.levels[-1] = 0
+        if last:
+            self._mark_last()
+        if child is None:
+            self._write_line(f"{label}: None")
         else:
-                indent += self.CHILD
+            self._write_line(label)
-        if self.idx is not None:
+            with self._child_level(last=True):
-            text = f"[{self.idx}] {text}"
+                child.accept(self)
-        self.idx = None
-        return indent + text + "\n"
 
-    def visit_type_expr(self, expr: TypeExpr) -> str:
+    def visit_type_expr(self, expr: TypeExpr):
-        res: str = self.print_line("TypeExpr")
+        self._write_line("TypeExpr")
-        self.levels.append(1)
+        with self._child_level():
-        res += self.print_line(f'name: "{expr.name.lexeme}"')
+            self._write_line(f'name: "{expr.name.lexeme}"')
-        self.levels[-1] = 2
+            self._write_optional_child("schema", expr.schema, last=True)
-        if expr.schema is None:
-            res += self.print_line("schema: None")
-        else:
-            res += self.print_line("schema")
-            self.levels.append(2)
-            res += expr.schema.accept(self)
-            self.levels.pop()
-        self.levels.pop()
-        return res
 
-    def visit_schema_expr(self, expr: SchemaExpr) -> str:
+    def visit_schema_expr(self, expr: SchemaExpr):
-        res: str = self.print_line("SchemaExpr")
+        self._write_line("SchemaExpr")
-        self.levels.append(1)
+        with self._child_level():
             for i, elmt in enumerate(expr.elements):
-            self.idx = i
+                self._idx = i
                 if i == len(expr.elements) - 1:
-                self.levels[-1] = 2
+                    self._mark_last()
-            res += elmt.accept(self)
+                elmt.accept(self)
-        self.levels.pop()
-        return res
 
-    def visit_schema_element_expr(self, expr: SchemaElementExpr) -> str:
+    def visit_schema_element_expr(self, expr: SchemaElementExpr):
-        res: str = self.print_line("SchemaElementExpr")
+        self._write_line("SchemaElementExpr")
-        self.levels.append(1)
+        with self._child_level():
-        res += self.print_line(
+            name_text: str = "None" if expr.name is None else f'"{expr.name.lexeme}"'
-            "name: " + ("None" if expr.name is None else f'"{expr.name.lexeme}"')
+            self._write_line(f"name: {name_text}")
-        )
+            self._write_optional_child("type", expr.type, last=True)
-        self.levels[-1] = 2
-        if expr.type is None:
-            res += self.print_line("type: None")
-        else:
-            res += self.print_line("type")
-            self.levels.append(2)
-            res += expr.type.accept(self)
-            self.levels.pop()
-        self.levels.pop()
-        return res
 
 
 class AnnotationPrinter(Expr.Visitor[str]):