How to Mock a Catalyst Logical Plan in Python

from typing import Any, Dict, List, Optional


class CatalystLogicalPlan:
    """A minimal mock of Catalyst's logical plan for teaching purposes."""
    
    def __init__(self, node_type: str, **kwargs: Any) -> None:
        self.node_type = node_type
        self.attributes: Dict[str, Any] = kwargs
        self.children: List["CatalystLogicalPlan"] = []

    def add_child(self, child: "CatalystLogicalPlan") -> None:
        self.children.append(child)

    def output_columns(self) -> List[str]:
        if "output" in self.attributes:
            return self.attributes["output"]
        if self.children:
            return self.children[-1].output_columns()
        return []

    def explain(self, indent: int = 0) -> str:
        prefix = "  " * indent
        details = ", ".join(f"{k}={v}" for k, v in self.attributes.items())
        result = [f"{prefix}+- {self.node_type} ({details})"]
        for child in self.children:
            result.append(child.explain(indent + 1))
        return "\n".join(result)


def build_mock_plan() -> CatalystLogicalPlan:
    scan = CatalystLogicalPlan("Project", output=["id", "name", "salary"])
    filter_node = CatalystLogicalPlan("Filter", predicate="salary > 50000")
    aggregate = CatalystLogicalPlan("Aggregate", grouping_keys=["department"], agg_func="avg(salary)")
    
    # Build tree: Aggregate <- Filter <- Scan
    scan.add_child(CatalystLogicalPlan("Scan", table="employees"))
    filter_node.add_child(scan)
    aggregate.add_child(filter_node)
    return aggregate


if __name__ == "__main__":
    plan = build_mock_plan()
    print(plan.explain())
    print(f"Output columns: {plan.output_columns()}")