| """Safe parser for the SRHarness symbolic model language.""" |
|
|
| from __future__ import annotations |
|
|
| import ast |
| from collections.abc import Mapping |
| from typing import Any |
|
|
| from .analysis import fold_constants |
| from .desugar import desugar |
| from .expression import ( |
| Aggregate, |
| Binary, |
| Expression, |
| Function, |
| Gather, |
| GroupedParameter, |
| Index, |
| Indexed, |
| Number, |
| Parameter, |
| Reduction, |
| RelationLift, |
| Symbol, |
| Unary, |
| ) |
|
|
|
|
| FUNCTIONS = { |
| "abs", "arccos", "arcsin", "arctan", "cos", "cosh", "cot", "csc", "exp", "inv", |
| "log", "log10", "log2", "sec", "sech", "sigmoid", "sign", "sin", "sinh", "sqrt", "tan", |
| "tanh", "pow2", "pow3", |
| } |
| FUNCTION_ALIASES = { |
| "acos": "arccos", |
| "asin": "arcsin", |
| "atan": "arctan", |
| "lg": "log10", |
| "ln": "log", |
| } |
| BINARY_FUNCTIONS = {"max", "min"} |
|
|
|
|
| class ExpressionParser(ast.NodeVisitor): |
| """Convert a restricted Python expression AST into engine nodes.""" |
|
|
| def __init__(self, symbols: Mapping[str, Any] | None = None): |
| self.symbols = {} |
| for name, value in (symbols or {}).items(): |
| if isinstance(value, Expression): |
| self.symbols[name] = value |
| elif isinstance(value, bool) or not isinstance(value, (int, float)): |
| raise TypeError(f"Symbol override {name!r} must be an expression or number.") |
| else: |
| self.symbols[name] = Number(value) |
|
|
| def parse(self, source: str) -> Expression: |
| """Parse input into the canonical representation. |
| |
| Args: |
| source: Source text to parse. |
| |
| Returns: |
| The parsed, desugared, and constant-folded expression. |
| """ |
| try: |
| tree = ast.parse(source, mode="eval") |
| except SyntaxError as error: |
| raise SyntaxError(f"Invalid symbolic expression: {error.msg}.") from error |
| return fold_constants(desugar(self.visit(tree.body))) |
|
|
| def generic_visit(self, node: ast.AST): |
| """Reject syntax that is outside the symbolic language. |
| |
| Args: |
| node: Expression or syntax-tree node. |
| """ |
| raise ValueError(f"Unsupported syntax: {type(node).__name__}.") |
|
|
| def visit_Constant(self, node: ast.Constant) -> Expression: |
| """Convert a numeric literal into a number node. |
| |
| Args: |
| node: Expression or syntax-tree node. |
| |
| Returns: |
| A numeric expression node. |
| """ |
| if isinstance(node.value, bool) or not isinstance(node.value, (int, float)): |
| raise ValueError(f"Only numerical literals are expressions, got {node.value!r}.") |
| return Number(node.value) |
|
|
| def visit_Name(self, node: ast.Name) -> Expression: |
| """Resolve a name to a supplied or newly created symbol. |
| |
| Args: |
| node: Expression or syntax-tree node. |
| |
| Returns: |
| The expression associated with the name. |
| """ |
| return self.symbols.setdefault(node.id, Symbol(node.id)) |
|
|
| def visit_UnaryOp(self, node: ast.UnaryOp) -> Expression: |
| """Convert a supported unary operator. |
| |
| Args: |
| node: Expression or syntax-tree node. |
| |
| Returns: |
| The converted operand or unary expression. |
| """ |
| if isinstance(node.op, ast.USub): |
| return Unary("-", self.visit(node.operand)) |
| if isinstance(node.op, ast.UAdd): |
| return self.visit(node.operand) |
| raise ValueError(f"Unsupported unary operator: {type(node.op).__name__}.") |
|
|
| def visit_BinOp(self, node: ast.BinOp) -> Expression: |
| """Convert a supported binary arithmetic operator. |
| |
| Args: |
| node: Expression or syntax-tree node. |
| |
| Returns: |
| A binary expression node. |
| """ |
| operators = { |
| ast.Add: "+", ast.Sub: "-", ast.Mult: "*", ast.Div: "/", ast.Pow: "**" |
| } |
| operator = operators.get(type(node.op)) |
| if operator is None: |
| raise ValueError(f"Unsupported binary operator: {type(node.op).__name__}.") |
| return Binary(operator, self.visit(node.left), self.visit(node.right)) |
|
|
| def visit_Subscript(self, node: ast.Subscript) -> Expression: |
| """Convert symbolic indexing such as ``x[i]``. |
| |
| Args: |
| node: Expression or syntax-tree node. |
| |
| Returns: |
| An indexed expression node. |
| """ |
| base = self.visit(node.value) |
| return Indexed(base, self._indices(node.slice)) |
|
|
| def visit_Call(self, node: ast.Call) -> Expression: |
| """Convert an approved symbolic function or reduction call. |
| |
| Args: |
| node: Expression or syntax-tree node. |
| |
| Returns: |
| The corresponding symbolic expression. |
| """ |
| if isinstance(node.func, ast.Subscript) and self._name(node.func.value) == "sum": |
| if node.keywords or len(node.args) not in {1, 2}: |
| raise ValueError( |
| "sum[index](...) expects an operand, optionally preceded by a relation." |
| ) |
| arguments = [self.visit(argument) for argument in node.args] |
| relation, operand = (None, arguments[0]) if len(arguments) == 1 else arguments |
| return Reduction(self._indices(node.func.slice), operand, relation) |
|
|
| source_name = self._name(node.func) |
| name = FUNCTION_ALIASES.get(source_name, source_name) |
| if name in FUNCTIONS: |
| if node.keywords or len(node.args) != 1: |
| raise ValueError(f"{source_name}(...) expects exactly one positional argument.") |
| return Function(name, (self.visit(node.args[0]),)) |
| if name in BINARY_FUNCTIONS: |
| if node.keywords or len(node.args) != 2: |
| raise ValueError(f"{name}(...) expects exactly two positional arguments.") |
| return Function(name, tuple(self.visit(argument) for argument in node.args)) |
| if name == "delay": |
| if node.keywords or len(node.args) != 2: |
| raise ValueError("delay(value, delta) expects exactly two arguments.") |
| return Function(name, tuple(self.visit(argument) for argument in node.args)) |
| if name == "gather": |
| if node.keywords or len(node.args) != 2: |
| raise ValueError("gather(relation, expression) expects exactly two arguments.") |
| return Gather(self.visit(node.args[0]), self.visit(node.args[1])) |
| if name == "Number": |
| if node.keywords or len(node.args) != 1: |
| raise ValueError("Number(...) expects exactly one numerical literal.") |
| return self.visit(node.args[0]) |
| if name == "param": |
| return self._parameter(node) |
| if name == "grouped_param": |
| return self._grouped_parameter(node) |
| if name == "aggr": |
| if node.keywords or len(node.args) not in {1, 2}: |
| raise ValueError( |
| "aggr expects aggr(relation * expression) or " |
| "aggr(relation, expression)." |
| ) |
| if len(node.args) == 2: |
| return Aggregate(self.visit(node.args[0]), self.visit(node.args[1])) |
| relation, operand = _split_aggregation(self.visit(node.args[0])) |
| return Aggregate(relation, operand) |
| if name in {"targ", "sour"}: |
| if node.keywords or len(node.args) not in {1, 2}: |
| raise ValueError(f"{name} expects one value, optionally preceded by a relation.") |
| arguments = [self.visit(argument) for argument in node.args] |
| relation, operand = (None, arguments[0]) if len(arguments) == 1 else arguments |
| return RelationLift("target" if name == "targ" else "source", operand, relation) |
| raise ValueError(f"Unknown symbolic function: {name!r}.") |
|
|
| def _parameter(self, node: ast.Call) -> Parameter: |
| if len(node.args) != 1 or not isinstance(node.args[0], ast.Constant): |
| raise ValueError("param(...) requires a literal string name.") |
| name = node.args[0].value |
| if not isinstance(name, str) or not name: |
| raise ValueError("Parameter names must be non-empty strings.") |
| keywords = self._literal_keywords(node, {"value"}) |
| value = keywords.get("value") |
| if value is not None and (isinstance(value, bool) or not isinstance(value, (int, float))): |
| raise ValueError("A parameter initial value must be numerical.") |
| return Parameter(name, value) |
|
|
| def _grouped_parameter(self, node: ast.Call) -> GroupedParameter: |
| if len(node.args) != 1: |
| raise ValueError("grouped_param(...) requires one grouping expression.") |
| keywords = self._literal_keywords(node, {"name", "value", "default"}) |
| name = keywords.get("name") |
| if name is not None and (not isinstance(name, str) or not name): |
| raise ValueError("A grouped parameter name must be a non-empty string.") |
| value = keywords.get("value") |
| if value is not None and not isinstance(value, dict): |
| raise ValueError("grouped_param value must be a category-to-number dictionary.") |
| default = keywords.get("default") |
| return GroupedParameter(self.visit(node.args[0]), name=name, value=value, default=default) |
|
|
| @staticmethod |
| def _name(node: ast.AST) -> str: |
| if not isinstance(node, ast.Name): |
| raise ValueError("Only direct symbolic function calls are allowed.") |
| return node.id |
|
|
| @staticmethod |
| def _indices(node: ast.AST) -> tuple[Index, ...]: |
| nodes = node.elts if isinstance(node, ast.Tuple) else (node,) |
| if not nodes or not all(isinstance(item, ast.Name) for item in nodes): |
| raise ValueError("Indices must be identifiers such as i, j, or k.") |
| return tuple(Index(item.id) for item in nodes) |
|
|
| @staticmethod |
| def _literal_keywords(node: ast.Call, allowed: set[str]) -> dict[str, Any]: |
| result = {} |
| for keyword in node.keywords: |
| if keyword.arg not in allowed: |
| raise ValueError(f"Unexpected keyword argument: {keyword.arg!r}.") |
| result[keyword.arg] = ast.literal_eval(keyword.value) |
| return result |
|
|
|
|
| def parse( |
| source: str, |
| symbols: Mapping[str, Any] | None = None, |
| ) -> Expression: |
| """Parse *source* without using ``eval`` or executing user code. |
| |
| Args: |
| source: Source text to parse. |
| symbols: Optional predefined symbols or numeric constants. |
| |
| Returns: |
| The parsed canonical expression. |
| """ |
| return ExpressionParser(symbols).parse(source) |
|
|
|
|
| def _split_aggregation(expression: Expression) -> tuple[Expression, Expression]: |
| factors = _multiplication_factors(expression) |
| if len(factors) < 2 or not isinstance(factors[0], Symbol): |
| raise ValueError( |
| "aggr(relation * expression) requires the relation symbol as its first factor." |
| ) |
| operand = factors[1] |
| for factor in factors[2:]: |
| operand = Binary("*", operand, factor) |
| return factors[0], operand |
|
|
|
|
| def _multiplication_factors(expression: Expression) -> list[Expression]: |
| if isinstance(expression, Binary) and expression.operator == "*": |
| return _multiplication_factors(expression.left) + _multiplication_factors(expression.right) |
| return [expression] |
|
|
