File size: 5,009 Bytes
d06fa57 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | """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))
|