vent_agent_2.0/algorithms/airflow_review/engine/expression.py

"""规则库内表达式的确定性求值器。

规则库把公式写成字符串，这里用白名单 AST 求值，既能保证「公式不由大模型
生成」，又能让公式与参数分离、可版本化。不引入 eval 的开放执行面。
"""
from __future__ import annotations

import ast
import math
from typing import Any, Mapping


class ExpressionError(ValueError):
    """表达式非法或引用了未提供的参数。"""


_ALLOWED_FUNCS = {
    "max": max,
    "min": min,
    "sum": sum,
    "abs": abs,
    "round": round,
    "ceil": math.ceil,
    "floor": math.floor,
    "len": len,
}
_ALLOWED_NODES = (
    ast.Expression, ast.BinOp, ast.UnaryOp, ast.Constant, ast.Name, ast.Load,
    ast.Call, ast.List, ast.Tuple, ast.Set, ast.Compare, ast.IfExp, ast.BoolOp,
    ast.Add, ast.Sub, ast.Mult, ast.Div, ast.Mod, ast.Pow, ast.USub, ast.UAdd,
    ast.Gt, ast.GtE, ast.Lt, ast.LtE, ast.Eq, ast.NotEq, ast.And, ast.Or, ast.Not,
    ast.LtE, ast.keyword,
)
# 表达式规模上限：规则库里的公式都是短式，超限说明输入异常
_MAX_LENGTH = 500
_MAX_POWER = 8


def evaluate(expression: str, params: Mapping[str, Any]) -> float:
    """按 params 求值；参数缺失抛 ExpressionError，由调用方转"依据不足"。"""
    if not isinstance(expression, str) or not expression.strip():
        raise ExpressionError("表达式为空")
    if len(expression) > _MAX_LENGTH:
        raise ExpressionError("表达式过长")
    try:
        tree = ast.parse(expression, mode="eval")
    except SyntaxError as exc:
        raise ExpressionError(f"表达式语法错误：{exc.msg}") from exc

    for node in ast.walk(tree):
        if not isinstance(node, _ALLOWED_NODES):
            raise ExpressionError(f"表达式含不允许的语法：{type(node).__name__}")
        if isinstance(node, ast.Call):
            if not isinstance(node.func, ast.Name) or node.func.id not in _ALLOWED_FUNCS:
                name = getattr(node.func, "id", type(node.func).__name__)
                raise ExpressionError(f"表达式含不允许的函数：{name}")
        if isinstance(node, ast.Constant) and isinstance(node.value, str):
            raise ExpressionError("表达式不允许字符串常量")
        if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Pow):
            exponent = getattr(node.right, "value", None)
            if not isinstance(exponent, (int, float)) or abs(exponent) > _MAX_POWER:
                raise ExpressionError("幂次超出允许范围")

    try:
        value = _eval_node(tree.body, params)
    except ExpressionError:
        raise
    except (TypeError, ValueError, ZeroDivisionError, KeyError) as exc:
        raise ExpressionError(f"表达式求值失败：{exc}") from exc
    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise ExpressionError("表达式未返回数值")
    return float(value)


def try_evaluate(expression: str, params: Mapping[str, Any]) -> tuple[float | None, str | None]:
    """求值失败不抛异常，返回 (值, 失败原因)，供"依据不足"分支使用。"""
    try:
        return evaluate(expression, params), None
    except ExpressionError as exc:
        return None, str(exc)


def missing_params(expression: str, params: Mapping[str, Any]) -> list[str]:
    """列出表达式引用但 params 中缺失的变量名（按在表达式中出现的先后顺序）。"""
    try:
        tree = ast.parse(expression, mode="eval")
    except SyntaxError:
        return []
    names: list[str] = []
    for node in sorted(ast.walk(tree), key=lambda item: (getattr(item, "lineno", 0),
                                                         getattr(item, "col_offset", 0))):
        if isinstance(node, ast.Name) and node.id not in _ALLOWED_FUNCS:
            if node.id not in params and node.id not in names:
                names.append(node.id)
    return names


def _eval_node(node: ast.AST, params: Mapping[str, Any]) -> Any:
    if isinstance(node, ast.Constant):
        return node.value
    if isinstance(node, ast.Name):
        if node.id in _ALLOWED_FUNCS:
            return _ALLOWED_FUNCS[node.id]
        if node.id not in params:
            raise ExpressionError(f"缺少参数：{node.id}")
        return params[node.id]
    if isinstance(node, ast.Call):
        func = _eval_node(node.func, params)
        if not callable(func):
            raise ExpressionError("调用目标不是函数")
        args = [_eval_node(arg, params) for arg in node.args]
        return func(*args)
    if isinstance(node, ast.List):
        return [_eval_node(item, params) for item in node.elts]
    if isinstance(node, ast.Tuple):
        return tuple(_eval_node(item, params) for item in node.elts)
    if isinstance(node, ast.Set):
        return {_eval_node(item, params) for item in node.elts}
    if isinstance(node, ast.UnaryOp):
        operand = _eval_node(node.operand, params)
        if isinstance(node.op, ast.USub):
            return -operand
        if isinstance(node.op, ast.UAdd):
            return +operand
        if isinstance(node.op, ast.Not):
            return not operand
        raise ExpressionError("不支持的一元运算符")
    if isinstance(node, ast.BinOp):
        left = _eval_node(node.left, params)
        right = _eval_node(node.right, params)
        return _apply_binop(node.op, left, right)
    if isinstance(node, ast.BoolOp):
        values = [_eval_node(value, params) for value in node.values]
        return all(values) if isinstance(node.op, ast.And) else any(values)
    if isinstance(node, ast.Compare):
        left = _eval_node(node.left, params)
        for op, comparator in zip(node.ops, node.comparators):
            right = _eval_node(comparator, params)
            if not _compare(op, left, right):
                return False
            left = right
        return True
    if isinstance(node, ast.IfExp):
        return _eval_node(node.body, params) if _eval_node(node.test, params) else _eval_node(node.orelse, params)
    raise ExpressionError(f"不支持的语法：{type(node).__name__}")


def _apply_binop(op: ast.AST, left: Any, right: Any) -> Any:
    if isinstance(op, ast.Add):
        return left + right
    if isinstance(op, ast.Sub):
        return left - right
    if isinstance(op, ast.Mult):
        return left * right
    if isinstance(op, ast.Div):
        return left / right
    if isinstance(op, ast.Mod):
        return left % right
    if isinstance(op, ast.Pow):
        return left ** right
    raise ExpressionError(f"不支持的运算符：{type(op).__name__}")


def _compare(op: ast.AST, left: Any, right: Any) -> bool:
    if isinstance(op, ast.Gt):
        return left > right
    if isinstance(op, ast.GtE):
        return left >= right
    if isinstance(op, ast.Lt):
        return left < right
    if isinstance(op, ast.LtE):
        return left <= right
    if isinstance(op, ast.Eq):
        return left == right
    if isinstance(op, ast.NotEq):
        return left != right
    raise ExpressionError(f"不支持的比较符：{type(op).__name__}")