| """Immutable expression-tree traversal and transformation helpers.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import replace as dataclass_replace |
|
|
| from .expression import ( |
| Aggregate, |
| Binary, |
| Expression, |
| Function, |
| Gather, |
| GroupedParameter, |
| Indexed, |
| Reduction, |
| RelationLift, |
| Unary, |
| ) |
|
|
|
|
| def children(node: Expression) -> tuple[Expression, ...]: |
| """Return the direct child expressions of a node. |
| |
| Args: |
| node: Expression or syntax-tree node. |
| |
| Returns: |
| The node's direct children in structural order. |
| """ |
| if isinstance(node, Unary): |
| return (node.operand,) |
| if isinstance(node, Binary): |
| return (node.left, node.right) |
| if isinstance(node, Function): |
| return node.arguments |
| if isinstance(node, Indexed): |
| return (node.base,) |
| if isinstance(node, Reduction): |
| return (node.operand,) if node.relation is None else (node.relation, node.operand) |
| if isinstance(node, Gather): |
| return (node.relation, node.operand) |
| if isinstance(node, Aggregate): |
| return (node.relation, node.operand) |
| if isinstance(node, RelationLift): |
| return (node.operand,) if node.relation is None else (node.relation, node.operand) |
| if isinstance(node, GroupedParameter): |
| return (node.by,) |
| return () |
|
|
|
|
| def with_children(node: Expression, values: tuple[Expression, ...]) -> Expression: |
| """Return a copy of a node with new direct children. |
| |
| Args: |
| node: Expression or syntax-tree node. |
| values: Replacement children in structural order. |
| |
| Returns: |
| The rebuilt expression node. |
| """ |
| if len(values) != len(children(node)): |
| raise ValueError(f"Expected {len(children(node))} children, got {len(values)}.") |
| if isinstance(node, Unary): |
| return dataclass_replace(node, operand=values[0]) |
| if isinstance(node, Binary): |
| return dataclass_replace(node, left=values[0], right=values[1]) |
| if isinstance(node, Function): |
| return dataclass_replace(node, arguments=values) |
| if isinstance(node, Indexed): |
| return dataclass_replace(node, base=values[0]) |
| if isinstance(node, Reduction): |
| if node.relation is None: |
| return dataclass_replace(node, operand=values[0]) |
| return dataclass_replace(node, relation=values[0], operand=values[1]) |
| if isinstance(node, Gather): |
| return dataclass_replace(node, relation=values[0], operand=values[1]) |
| if isinstance(node, Aggregate): |
| return dataclass_replace(node, relation=values[0], operand=values[1]) |
| if isinstance(node, RelationLift): |
| if node.relation is None: |
| return dataclass_replace(node, operand=values[0]) |
| return dataclass_replace(node, relation=values[0], operand=values[1]) |
| if isinstance(node, GroupedParameter): |
| return dataclass_replace(node, by=values[0]) |
| return node |
|
|
|
|
| def iter_preorder(node: Expression): |
| """Yield an expression tree in preorder. |
| |
| Args: |
| node: Expression or syntax-tree node. |
| |
| Yields: |
| The node followed recursively by its children. |
| """ |
| yield node |
| for child in children(node): |
| yield from iter_preorder(child) |
|
|
|
|
| def iter_postorder(node: Expression): |
| """Yield an expression tree in postorder. |
| |
| Args: |
| node: Expression or syntax-tree node. |
| |
| Yields: |
| Descendants followed by their parent node. |
| """ |
| for child in children(node): |
| yield from iter_postorder(child) |
| yield node |
|
|
|
|
| def transform(node: Expression, function) -> Expression: |
| """Apply a bottom-up transformation to an expression tree. |
| |
| Args: |
| node: Expression or syntax-tree node. |
| function: Transformation callback. |
| |
| Returns: |
| The transformed expression tree. |
| """ |
| transformed = with_children(node, tuple(transform(child, function) for child in children(node))) |
| return function(transformed) |
|
|
|
|
| def replace(node: Expression, old: Expression, new: Expression) -> Expression: |
| """Replace a node by object identity throughout a tree. |
| |
| Args: |
| node: Expression or syntax-tree node. |
| old: Existing expression node to replace. |
| new: Replacement expression node. |
| |
| Returns: |
| The rebuilt expression tree. |
| """ |
| if node is old: |
| return new |
| return with_children(node, tuple(replace(child, old, new) for child in children(node))) |
|
|
|
|
| def replace_at_path(node: Expression, path: tuple[int, ...], new: Expression) -> Expression: |
| """Replace the node at a tuple of child indices. |
| |
| Args: |
| node: Expression or syntax-tree node. |
| path: Child indices from the root to the target node. |
| new: Replacement expression node. |
| |
| Returns: |
| The rebuilt expression tree. |
| """ |
| if not path: |
| return new |
| values = list(children(node)) |
| index = path[0] |
| values[index] = replace_at_path(values[index], path[1:], new) |
| return with_children(node, tuple(values)) |
|
|
