zkxw/backend/app/services/oa_formula.py

"""OA 表单公式与数据联动（禁止任意 eval）"""
from __future__ import annotations

import ast
import operator
import re
from typing import Any, Dict, List, Optional, Tuple, Union

Number = Union[int, float]

_BIN_OPS = {
    ast.Add: operator.add,
    ast.Sub: operator.sub,
    ast.Mult: operator.mul,
    ast.Div: operator.truediv,
    ast.Mod: operator.mod,
}

_CMP_OPS = {
    ast.Eq: operator.eq,
    ast.NotEq: operator.ne,
    ast.Lt: operator.lt,
    ast.LtE: operator.le,
    ast.Gt: operator.gt,
    ast.GtE: operator.ge,
}

# 字段引用：{field} 或 field（Ident）
_FIELD_BRACE = re.compile(r"\{([A-Za-z_][A-Za-z0-9_]*)\}")


class FormulaError(ValueError):
    """公式非法或求值失败"""


def _to_number(v: Any) -> Number:
    if v is None or v == "":
        return 0
    if isinstance(v, bool):
        return int(v)
    if isinstance(v, (int, float)):
        return v
    if isinstance(v, str):
        s = v.strip()
        if not s:
            return 0
        try:
            if "." in s:
                return float(s)
            return int(s)
        except ValueError as e:
            raise FormulaError(f"无法将「{v}」转为数字") from e
    raise FormulaError(f"不支持的数值类型: {type(v).__name__}")


def _truthy(v: Any) -> bool:
    if v is None:
        return False
    if isinstance(v, (int, float)):
        return v != 0
    if isinstance(v, str):
        return v.strip() != ""
    if isinstance(v, (list, dict)):
        return len(v) > 0
    return bool(v)


def preprocess_formula(expr: str) -> str:
    """将 {field} 转为合法标识符引用 field。"""
    text = (expr or "").strip()
    if not text:
        raise FormulaError("公式为空")
    return _FIELD_BRACE.sub(r"\1", text)


class _SafeEval(ast.NodeVisitor):
    def __init__(self, form_data: Dict[str, Any]):
        self.form_data = form_data or {}

    def visit(self, node):  # type: ignore[override]
        method = "visit_" + node.__class__.__name__
        visitor = getattr(self, method, None)
        if visitor is None:
            raise FormulaError(f"不支持的语法: {node.__class__.__name__}")
        return visitor(node)

    def visit_Expression(self, node: ast.Expression):
        return self.visit(node.body)

    def visit_Constant(self, node: ast.Constant):
        if isinstance(node.value, (int, float, str, bool)) or node.value is None:
            return node.value
        raise FormulaError("不支持的常量类型")

    # Py3.9 可能仍有 Num/Str/NameConstant
    def visit_Num(self, node):  # pragma: no cover
        return node.n

    def visit_Str(self, node):  # pragma: no cover
        return node.s

    def visit_NameConstant(self, node):  # pragma: no cover
        return node.value

    def visit_Name(self, node: ast.Name):
        if not isinstance(node.ctx, ast.Load):
            raise FormulaError("仅允许读取字段")
        key = node.id
        if key not in self.form_data:
            # 未填视为 0 / 空：数字上下文由运算再转
            return self.form_data.get(key)
        return self.form_data.get(key)

    def visit_UnaryOp(self, node: ast.UnaryOp):
        val = self.visit(node.operand)
        if isinstance(node.op, ast.UAdd):
            return +_to_number(val)
        if isinstance(node.op, ast.USub):
            return -_to_number(val)
        if isinstance(node.op, ast.Not):
            return not _truthy(val)
        raise FormulaError("不支持的一元运算")

    def visit_BinOp(self, node: ast.BinOp):
        op_type = type(node.op)
        if op_type not in _BIN_OPS:
            raise FormulaError("仅支持 + - * / %")
        left = _to_number(self.visit(node.left))
        right = _to_number(self.visit(node.right))
        if op_type is ast.Div and right == 0:
            raise FormulaError("除数不能为 0")
        if op_type is ast.Mod and right == 0:
            raise FormulaError("模运算除数不能为 0")
        return _BIN_OPS[op_type](left, right)

    def visit_Compare(self, node: ast.Compare):
        left = self.visit(node.left)
        for op, comp in zip(node.ops, node.comparators):
            op_type = type(op)
            if op_type not in _CMP_OPS:
                raise FormulaError("不支持的比较运算")
            right = self.visit(comp)
            # 两侧能数字则按数字比
            try:
                l2, r2 = _to_number(left), _to_number(right)
                ok = _CMP_OPS[op_type](l2, r2)
            except FormulaError:
                ok = _CMP_OPS[op_type](left, right)
            if not ok:
                return False
            left = right
        return True

    def visit_BoolOp(self, node: ast.BoolOp):
        if isinstance(node.op, ast.And):
            val = True
            for v in node.values:
                val = _truthy(self.visit(v))
                if not val:
                    return False
            return True
        if isinstance(node.op, ast.Or):
            for v in node.values:
                if _truthy(self.visit(v)):
                    return True
            return False
        raise FormulaError("不支持的逻辑运算")

    def visit_IfExp(self, node: ast.IfExp):
        return self.visit(node.body) if _truthy(self.visit(node.test)) else self.visit(node.orelse)

    def visit_Call(self, node: ast.Call):
        if not isinstance(node.func, ast.Name):
            raise FormulaError("仅支持白名单函数")
        fname = node.func.id.upper()
        if node.keywords:
            raise FormulaError("函数不支持关键字参数")
        args = [self.visit(a) for a in node.args]
        if fname == "IF":
            if len(args) != 3:
                raise FormulaError("IF 需要 3 个参数")
            return args[1] if _truthy(args[0]) else args[2]
        if fname == "SUM":
            total: Number = 0
            for a in args:
                if isinstance(a, list):
                    for item in a:
                        total = total + _to_number(item)
                else:
                    total = total + _to_number(a)
            return total
        if fname == "ABS":
            if len(args) != 1:
                raise FormulaError("ABS 需要 1 个参数")
            return abs(_to_number(args[0]))
        raise FormulaError(f"函数未开放: {node.func.id}")

    def visit_List(self, node: ast.List):
        return [self.visit(e) for e in node.elts]


def eval_formula(expr: str, form_data: Optional[Dict[str, Any]] = None) -> Any:
    text = preprocess_formula(expr)
    try:
        tree = ast.parse(text, mode="eval")
    except SyntaxError as e:
        raise FormulaError(f"公式语法错误: {e.msg}") from e
    return _SafeEval(form_data or {}).visit(tree)


def _extract_rules(form_schema: Any) -> List[Dict[str, Any]]:
    if isinstance(form_schema, list):
        return [r for r in form_schema if isinstance(r, dict)]
    if isinstance(form_schema, dict):
        rules = form_schema.get("rule") or form_schema.get("rules") or form_schema.get("fields")
        if isinstance(rules, list):
            return [r for r in rules if isinstance(r, dict)]
    return []


def _walk_rules(rules: List[Dict[str, Any]], out: List[Dict[str, Any]]) -> None:
    for rule in rules:
        out.append(rule)
        children = rule.get("children") or []
        if isinstance(children, list):
            _walk_rules([c for c in children if isinstance(c, dict)], out)


def apply_formulas_and_linkages(
    schema: Any, form_data: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
    """
    按 schema 中字段 props.formula / props.linkage 回填 form_data。
    linkage 示例: { "when": "{days} > 3", "set": { "needGm": true } }
    或列表 form。
    """
    data: Dict[str, Any] = dict(form_data or {})
    flat: List[Dict[str, Any]] = []
    _walk_rules(_extract_rules(schema), flat)

    # 多轮公式：处理依赖（有上限）
    for _ in range(5):
        changed = False
        for rule in flat:
            field = rule.get("field")
            if not field:
                continue
            props = rule.get("props") or {}
            formula = props.get("formula") or rule.get("formula")
            if formula:
                try:
                    val = eval_formula(str(formula), data)
                except FormulaError:
                    raise
                if data.get(field) != val:
                    data[field] = val
                    changed = True
        if not changed:
            break

    for rule in flat:
        props = rule.get("props") or {}
        linkages = props.get("linkage") or rule.get("linkage")
        if not linkages:
            continue
        items = linkages if isinstance(linkages, list) else [linkages]
        for link in items:
            if not isinstance(link, dict):
                continue
            when = link.get("when") or link.get("condition")
            sets = link.get("set") or link.get("assign") or {}
            if when is None:
                continue
            try:
                ok = _truthy(eval_formula(str(when), data))
            except FormulaError:
                continue
            if ok and isinstance(sets, dict):
                for k, v in sets.items():
                    data[k] = v
    return data