237 lines
7.7 KiB
Python
237 lines
7.7 KiB
Python
"""受限的余额表达式引擎。
|
||
|
||
在 JSON 提取路径基础上支持：
|
||
- 运算符：+ - * / // % ** 一元负号、括号
|
||
- 函数（白名单）：float / int / abs / round / min / max / sum / len
|
||
- 路径写法与原来完全兼容：data.balance、data[0].x、$.data[0].x
|
||
|
||
示例：
|
||
data.balance / 100 # 分转元
|
||
float(data.balance) # 字符串转数字
|
||
data.granted + data.topped_up # 多字段求和
|
||
sum(data[0].balances) # 数组求和
|
||
round(data.balance, 2) * 0.9
|
||
|
||
安全性：无 eval、无任意变量访问，函数名严格白名单，路径只支持
|
||
dict 键 / list 索引访问。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
|
||
_PATH_SEG_RE = re.compile(r"([^.\[]+)|\[(-?\d+)\]")
|
||
|
||
|
||
class ExprError(ValueError):
|
||
"""表达式语法错误 / 求值错误。"""
|
||
|
||
|
||
# ---------- 路径提取（返回原始值，不强制数字） ----------
|
||
|
||
def extract_path(data: object, path: str) -> object:
|
||
p = path.strip()
|
||
if p.startswith("$"):
|
||
p = p[1:]
|
||
cur: object = data
|
||
for key, idx in _PATH_SEG_RE.findall(p):
|
||
if key:
|
||
if not isinstance(cur, dict) or key not in cur:
|
||
raise ExprError(f"路径不存在: {path}（在 {key!r} 处）")
|
||
cur = cur[key]
|
||
if idx != "":
|
||
n = int(idx)
|
||
if not isinstance(cur, list) or n >= len(cur):
|
||
raise ExprError(f"数组索引越界: {path}（索引 {n}）")
|
||
cur = cur[n]
|
||
return cur
|
||
|
||
|
||
def to_number(value: object, src: str) -> float:
|
||
if isinstance(value, bool) or not isinstance(value, (int, float, str)):
|
||
raise ExprError(f"余额不是数字: {src} -> {value!r}")
|
||
try:
|
||
return float(value)
|
||
except (TypeError, ValueError):
|
||
raise ExprError(f"余额不是数字: {src} -> {value!r}")
|
||
|
||
|
||
# ---------- 词法 ----------
|
||
|
||
_TOKEN_RE = re.compile(r"""
|
||
(?P<num>\d+(?:\.\d+)?)
|
||
| (?P<path>\$?\.?[A-Za-z_]\w*(?:\.[A-Za-z_]\w*|\[-?\d+\])*)
|
||
| (?P<op>//|\*\*|[+\-*/%(),])
|
||
| (?P<ws>\s+)
|
||
""", re.VERBOSE)
|
||
|
||
|
||
def _tokenize(text: str) -> list[tuple[str, str]]:
|
||
tokens: list[tuple[str, str]] = []
|
||
pos = 0
|
||
for m in _TOKEN_RE.finditer(text):
|
||
if m.start() != pos:
|
||
raise ExprError(f"表达式语法错误（第 {pos + 1} 字符附近）: {text!r}")
|
||
pos = m.end()
|
||
kind = m.lastgroup
|
||
if kind == "ws":
|
||
continue
|
||
val = m.group()
|
||
if kind == "path" and text[m.end():m.end() + 1] == "(" and "." not in val and "[" not in val:
|
||
tokens.append(("FUNC", val))
|
||
else:
|
||
tokens.append((kind.upper(), val))
|
||
if pos != len(text):
|
||
raise ExprError(f"表达式语法错误（第 {pos + 1} 字符附近）: {text!r}")
|
||
tokens.append(("EOF", ""))
|
||
return tokens
|
||
|
||
|
||
# ---------- 函数白名单 ----------
|
||
|
||
def _call(name: str, args: list) -> object:
|
||
if name == "float":
|
||
_need_arity(name, args, 1)
|
||
return to_number(args[0], name)
|
||
if name == "int":
|
||
_need_arity(name, args, 1)
|
||
return int(to_number(args[0], name))
|
||
if name == "abs":
|
||
_need_arity(name, args, 1)
|
||
return abs(to_number(args[0], name))
|
||
if name == "round":
|
||
if len(args) == 1:
|
||
return round(to_number(args[0], name))
|
||
if len(args) == 2:
|
||
return round(to_number(args[0], name), int(args[1]))
|
||
raise ExprError("round() 需要 1 或 2 个参数")
|
||
if name == "len":
|
||
_need_arity(name, args, 1)
|
||
if not isinstance(args[0], (list, dict, str)):
|
||
raise ExprError("len() 参数必须是数组/对象/字符串")
|
||
return len(args[0])
|
||
if name == "sum":
|
||
if len(args) == 1 and isinstance(args[0], list):
|
||
return sum(to_number(v, name) for v in args[0])
|
||
if not args:
|
||
raise ExprError("sum() 至少需要 1 个参数")
|
||
return sum(to_number(v, name) for v in args)
|
||
if name == "min":
|
||
values = args[0] if len(args) == 1 and isinstance(args[0], list) else args
|
||
if not values:
|
||
raise ExprError("min() 参数不能为空")
|
||
return min(to_number(v, name) for v in values)
|
||
if name == "max":
|
||
values = args[0] if len(args) == 1 and isinstance(args[0], list) else args
|
||
if not values:
|
||
raise ExprError("max() 参数不能为空")
|
||
return max(to_number(v, name) for v in values)
|
||
raise ExprError(f"不支持的函数: {name}()（可用: float/int/abs/round/min/max/sum/len）")
|
||
|
||
|
||
def _need_arity(name: str, args: list, n: int) -> None:
|
||
if len(args) != n:
|
||
raise ExprError(f"{name}() 需要 {n} 个参数，实际 {len(args)} 个")
|
||
|
||
|
||
# ---------- 语法分析 + 求值（递归下降，直接求值） ----------
|
||
|
||
class _Evaluator:
|
||
def __init__(self, data: object, text: str) -> None:
|
||
self.data = data
|
||
self.tokens = _tokenize(text)
|
||
self.pos = 0
|
||
|
||
def _peek(self) -> tuple[str, str]:
|
||
return self.tokens[self.pos]
|
||
|
||
def _next(self) -> tuple[str, str]:
|
||
tok = self.tokens[self.pos]
|
||
self.pos += 1
|
||
return tok
|
||
|
||
def _accept_op(self, *ops: str) -> tuple[str, str] | None:
|
||
kind, val = self._peek()
|
||
if kind == "OP" and val in ops:
|
||
return self._next()
|
||
return None
|
||
|
||
def _expect_op(self, op: str) -> None:
|
||
kind, val = self._next()
|
||
if kind != "OP" or val != op:
|
||
raise ExprError(f"期望 {op!r}，实际 {val!r}")
|
||
|
||
def evaluate(self) -> float:
|
||
value = self._expr()
|
||
if self._peek() != ("EOF", ""):
|
||
raise ExprError(f"表达式多余内容: {self._peek()!r}")
|
||
return to_number(value, "表达式")
|
||
|
||
def _expr(self) -> object:
|
||
v = self._term()
|
||
while (op := self._accept_op("+", "-")) is not None:
|
||
rhs = self._term()
|
||
v = v + rhs if op[1] == "+" else v - rhs
|
||
return v
|
||
|
||
def _term(self) -> object:
|
||
v = self._factor()
|
||
while (op := self._accept_op("*", "/", "//", "%")) is not None:
|
||
rhs = self._factor()
|
||
if op[1] == "*":
|
||
v = v * rhs
|
||
elif op[1] == "/":
|
||
v = v / rhs
|
||
elif op[1] == "//":
|
||
v = v // rhs
|
||
else:
|
||
v = v % rhs
|
||
return v
|
||
|
||
def _factor(self) -> object:
|
||
v = self._unary()
|
||
if self._accept_op("**") is not None:
|
||
rhs = self._factor() # 右结合
|
||
v = v ** rhs
|
||
return v
|
||
|
||
def _unary(self) -> object:
|
||
if (op := self._accept_op("-", "+")) is not None:
|
||
v = self._unary()
|
||
return -v if op[1] == "-" else v
|
||
return self._primary()
|
||
|
||
def _primary(self) -> object:
|
||
kind, val = self._peek()
|
||
if kind == "NUM":
|
||
self._next()
|
||
return float(val) if "." in val else int(val)
|
||
if kind == "PATH":
|
||
self._next()
|
||
return extract_path(self.data, val)
|
||
if kind == "FUNC":
|
||
return self._call()
|
||
if kind == "OP" and val == "(":
|
||
self._next()
|
||
v = self._expr()
|
||
self._expect_op(")")
|
||
return v
|
||
raise ExprError(f"表达式语法错误: 意外的 {kind} {val!r}")
|
||
|
||
def _call(self) -> object:
|
||
name = self._next()[1]
|
||
self._expect_op("(")
|
||
args: list = []
|
||
if not (self._peek()[0] == "OP" and self._peek()[1] == ")"):
|
||
while True:
|
||
args.append(self._expr())
|
||
if self._accept_op(",") is None:
|
||
break
|
||
self._expect_op(")")
|
||
return _call(name, args)
|
||
|
||
|
||
def evaluate_balance(data: object, expression: str) -> float:
|
||
"""对响应 JSON 求值余额表达式，返回 float。失败抛 ExprError。"""
|
||
return _Evaluator(data, expression).evaluate()
|