ast_nodes.py

"""
ast_nodes.py  –  ASTNode definition and pretty-printer.

Every parser rule returns an ASTNode (or a plain list for repeated items).
Nodes carry:
  - node_type : str   descriptive label, e.g. "ClassDecl", "MethodDecl"
  - children  : list  child ASTNodes (structural sub-trees)
  - value     : any   leaf data – identifiers, numbers, type names, etc.
  - line      : int   source line where this construct began (may be None)
"""


class ASTNode:
    __slots__ = ("node_type", "children", "value", "line")

    def __init__(self, node_type, children=None, value=None, line=None):
        self.node_type = node_type
        self.children  = children if children is not None else []
        self.value     = value
        self.line      = line

    # ------------------------------------------------------------------
    # Pretty-printer  (indented tree, similar to what IDEs show)
    # ------------------------------------------------------------------
    def pretty(self, indent: int = 0, last: bool = True, prefix: str = "") -> str:
        connector = "└── " if last else "├── "
        child_prefix = prefix + ("    " if last else "│   ")

        label = self.node_type
        if self.value is not None:
            label += f"  [{self.value}]"
        if self.line is not None:
            label += f"  (line {self.line})"

        lines = [prefix + connector + label]

        for i, child in enumerate(self.children):
            is_last = (i == len(self.children) - 1)
            if isinstance(child, ASTNode):
                lines.append(child.pretty(indent, is_last, child_prefix))
            elif isinstance(child, list):
                # Flatten plain lists (e.g. param lists) inline
                for j, item in enumerate(child):
                    sub_last = is_last and (j == len(child) - 1)
                    if isinstance(item, ASTNode):
                        lines.append(item.pretty(indent, sub_last, child_prefix))
                    else:
                        connector2 = "└── " if sub_last else "├── "
                        lines.append(child_prefix + connector2 + str(item))
            else:
                connector2 = "└── " if is_last else "├── "
                lines.append(child_prefix + connector2 + str(child))

        return "\n".join(lines)

    def __repr__(self):
        return f"ASTNode({self.node_type!r}, value={self.value!r}, line={self.line})"