feat(parser): use AST printer in test script

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/annotations.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Generic, Optional, TypeVar
+
+from lexer.token import Token
+
+T = TypeVar("T")
+
+
+@dataclass(frozen=True)
+class Expr(ABC):
+    @abstractmethod
+    def accept(self, visitor: Visitor[T]) -> T: ...
+
+    class Visitor(ABC, Generic[T]):
+        @abstractmethod
+        def visit_type_expr(self, expr: TypeExpr) -> T: ...
+
+        @abstractmethod
+        def visit_schema_expr(self, expr: SchemaExpr) -> T: ...
+
+        @abstractmethod
+        def visit_schema_element_expr(self, expr: SchemaElementExpr) -> T: ...
+
+
+@dataclass(frozen=True)
+class TypeExpr(Expr):
+    name: Token
+    schema: Optional[SchemaExpr]
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_type_expr(self)
+
+
+@dataclass(frozen=True)
+class SchemaExpr(Expr):
+    left: Token
+    elements: list[Expr]
+    right: Token
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_schema_expr(self)
+
+
+@dataclass(frozen=True)
+class SchemaElementExpr(Expr):
+    name: Optional[Token]
+    type: Optional[Expr]
+
+    def accept(self, visitor: Expr.Visitor[T]) -> T:
+        return visitor.visit_schema_element_expr(self)

core/ast/printer.py

@@ -0,0 +1,122 @@
+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 _Level(Enum):
+    EMPTY = auto()
+    ACTIVE = auto()
+    LAST = auto()
+
+
+class AnnotationAstPrinter(Expr.Visitor[None]):
+    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: Expr):
+        self._buf = io.StringIO()
+        expr.accept(self)
+        return self._buf.getvalue()
+
+    @contextmanager
+    def _child_level(self, last: bool = False) -> Generator[None, None, None]:
+        self._levels.append(_Level.LAST if last 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):
+        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[Expr], *, 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(last=True):
+                child.accept(self)
+
+    def visit_type_expr(self, expr: TypeExpr):
+        self._write_line("TypeExpr")
+        with self._child_level():
+            self._write_line(f'name: "{expr.name.lexeme}"')
+            self._write_optional_child("schema", expr.schema, last=True)
+
+    def visit_schema_expr(self, expr: SchemaExpr):
+        self._write_line("SchemaExpr")
+        with self._child_level():
+            for i, elmt in enumerate(expr.elements):
+                self._idx = i
+                if i == len(expr.elements) - 1:
+                    self._mark_last()
+                elmt.accept(self)
+
+    def visit_schema_element_expr(self, expr: SchemaElementExpr):
+        self._write_line("SchemaElementExpr")
+        with self._child_level():
+            name_text: str = "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)
+
+
+class AnnotationPrinter(Expr.Visitor[str]):
+    def print(self, expr: Expr):
+        return expr.accept(self)
+
+    def visit_type_expr(self, expr: TypeExpr) -> str:
+        schema: str = ""
+        if expr.schema is not None:
+            schema = expr.schema.accept(self)
+        return f"{expr.name.lexeme}{schema}"
+
+    def visit_schema_expr(self, expr: SchemaExpr) -> str:
+        res: str = expr.left.lexeme
+        res += ", ".join(elmt.accept(self) for elmt in expr.elements)
+        res += expr.right.lexeme
+        return res
+
+    def visit_schema_element_expr(self, expr: SchemaElementExpr) -> str:
+        parts: list[str] = []
+        if expr.name is not None:
+            parts.append(expr.name.lexeme)
+
+        if expr.type is None:
+            parts.append("_")
+        else:
+            parts.append(expr.type.accept(self))
+        return ": ".join(parts)

parser/annotations.py

@@ -0,0 +1,64 @@
+from typing import Optional
+
+from core.ast.annotations import Expr, SchemaElementExpr, SchemaExpr, TypeExpr
+from lexer.token import Token, TokenType
+from parser.base import Parser
+from parser.errors import ParsingError
+
+
+class AnnotationParser(Parser):
+    SYNC_BOUNDARY: set[TokenType] = set()
+
+    def parse(self) -> Optional[Expr]:
+        expression: Optional[Expr] = self.annotation()
+        if not self.is_at_end():
+            self.error(self.peek(), "Extra tokens")
+        return expression
+
+    def synchronize(self):
+        self.advance()
+        while not self.is_at_end():
+            if self.peek().type in self.SYNC_BOUNDARY:
+                return
+            self.advance()
+
+    def annotation(self) -> Optional[Expr]:
+        try:
+            return self.type()
+        except ParsingError:
+            self.synchronize()
+            return None
+
+    def type(self) -> TypeExpr:
+        name: Token = self.consume(TokenType.IDENTIFIER, "Expected type identifier")
+        schema: Optional[SchemaExpr] = None
+        if self.match(TokenType.LEFT_BRACKET):
+            schema = self.schema()
+        return TypeExpr(name=name, schema=schema)
+
+    def schema(self) -> SchemaExpr:
+        left: Token = self.previous()
+        elements: list[Expr] = []
+        while not self.check(TokenType.RIGHT_BRACKET) and not self.is_at_end():
+            elements.append(self.schema_element())
+            if not self.check(TokenType.RIGHT_BRACKET):
+                self.consume(TokenType.COMMA, "Expected ',' between schema elements")
+
+        right: Token = self.consume(TokenType.RIGHT_BRACKET, "Unclosed schema")
+        return SchemaExpr(left=left, elements=elements, right=right)
+
+    def schema_element(self) -> Expr:
+        if self.match(TokenType.UNDERSCORE):
+            return SchemaElementExpr(name=None, type=None)
+
+        if not self.check(TokenType.IDENTIFIER):
+            raise self.error(self.peek(), "Expected schema element")
+
+        name: Optional[Token] = None
+        type: Optional[TypeExpr] = None
+        if self.check_next(TokenType.COLON):
+            name = self.advance()
+            self.advance()
+        if not self.match(TokenType.UNDERSCORE):
+            type = self.type()
+        return SchemaElementExpr(name=name, type=type)

parser/base.py

@@ -54,7 +54,7 @@ class Parser(ABC, Generic[T]):
         )
         self.current: int = 0
         self.length: int = len(self.tokens)
-        self.errors: list[TokenError]
+        self.errors: list[TokenError] = []
 
     def error(self, token: Token, message: str):
         """Record an error
@@ -120,6 +120,26 @@ class Parser(ABC, Generic[T]):
             return False
         return self.peek().type == token_type
 
+    def check_next(self, token_type: TokenType) -> bool:
+        """Check whether the next token is of the given type
+
+        This function always returns False if the parser is at the EOF token
+
+        Args:
+            token_type (TokenType): the type of token to check
+
+        Returns:
+            bool: True if the current token is of the given type and not EOF
+        """
+        if self.is_at_end():
+            return False
+        if self.current + 1 >= self.length:
+            return False
+        token: Token = self.tokens[self.current + 1]
+        if token.type == TokenType.EOF:
+            return False
+        return token.type == token_type
+
     def advance(self) -> Token:
         """Consume and return the current token, if not at the EOF
 

test.py

@@ -1,10 +1,11 @@
 import importlib
 from pathlib import Path
 
+from core.ast.printer import AnnotationAstPrinter
 from lexer.annotations import AnnotationLexer
 from lexer.midas import MidasLexer
 from lexer.token import Token
-
+from parser.annotations import AnnotationParser
 
 # Frame annotation
 mod = importlib.import_module("examples.00_syntax_prototype.01_simple_types")
@@ -12,17 +13,20 @@ mod = importlib.import_module("examples.00_syntax_prototype.01_simple_types")
 annotation: str = mod.__annotations__["df"]
 lexer: AnnotationLexer = AnnotationLexer(annotation, "01_simple_types.py")
 tokens: list[Token] = lexer.process()
-print([
-    f"{t.type.name}('{t.lexeme}')"
-    for t in tokens
-])
+# print([f"{t.type.name}('{t.lexeme}')" for t in tokens])
+
+parser = AnnotationParser(tokens)
+parsed = parser.parse()
+print(parsed)
+for err in parser.errors:
+    print(err.get_report())
+printer = AnnotationAstPrinter()
+if parsed is not None:
+    print(printer.print(parsed))
 
 # Midas type definitions
 path: Path = Path("examples") / "00_syntax_prototype" / "02_custom_types.midas"
 definitions: str = path.read_text()
 midas_lexer: MidasLexer = MidasLexer(definitions, path.name)
 tokens = midas_lexer.process()
-print([
-    f"{t.type.name}('{t.lexeme}')"
-    for t in tokens
-])
+# print([f"{t.type.name}('{t.lexeme}')" for t in tokens])