beam: add safe arithmetic expression evaluator with tests
This commit is contained in:
beam authored and 
Volt committed
1 parent
219c04d004
commit
239767a176
2 files changed
+268
No files matched your search
@@ -0,0 +1,224 @@
|
||||
"""Safe arithmetic expression evaluator.
|
||||
|
||||
A recursive-descent parser that evaluates arithmetic expressions without
|
||||
ever calling eval/exec. Supports +, -, *, /, //, %, **, parentheses,
|
||||
unary minus/plus, constants (pi, e), and a whitelist of math functions.
|
||||
|
||||
Public API:
|
||||
evaluate(expr: str) -> float | int
|
||||
Raises ValueError on malformed input or disallowed names.
|
||||
|
||||
Designed for untrusted input: no attribute access, no calls except a
|
||||
fixed function whitelist, no name lookups beyond constants.
|
||||
"""
|
||||
|
||||
import math
|
||||
import re
|
||||
|
||||
# Token kinds
|
||||
_NUMBER = "NUMBER"
|
||||
_IDENT = "IDENT"
|
||||
_OP = "OP"
|
||||
_LPAREN = "LPAREN"
|
||||
_RPAREN = "RPAREN"
|
||||
_COMMA = "COMMA"
|
||||
_EOF = "EOF"
|
||||
|
||||
# Allowed functions (name -> callable). All take floats and return floats
|
||||
# except min/max/round which are variadic.
|
||||
_FUNCTIONS = {
|
||||
"abs": abs,
|
||||
"round": round,
|
||||
"min": min,
|
||||
"max": max,
|
||||
"sin": math.sin,
|
||||
"cos": math.cos,
|
||||
"tan": math.tan,
|
||||
"sqrt": math.sqrt,
|
||||
"log": math.log,
|
||||
"log10": math.log10,
|
||||
"exp": math.exp,
|
||||
"floor": math.floor,
|
||||
"ceil": math.ceil,
|
||||
}
|
||||
|
||||
_CONSTANTS = {
|
||||
"pi": math.pi,
|
||||
"e": math.e,
|
||||
}
|
||||
|
||||
# Multi-char operators must be tried before single-char ones.
|
||||
_OPERATORS = ["**", "//", "+", "-", "*", "/", "%", "(", ")"]
|
||||
|
||||
_TOKEN_RE = re.compile(
|
||||
r"\s*(?:(?P<num>\d+(?:\.\d+)?(?:[eE][+-]?\d+)?)"
|
||||
r"|(?P<ident>[A-Za-z_][A-Za-z0-9_]*)"
|
||||
r"|(?P<op>\*\*|//|[+\-*/%(),]))"
|
||||
)
|
||||
|
||||
|
||||
class _Tokenizer:
|
||||
def __init__(self, text):
|
||||
self.text = text
|
||||
self.pos = 0
|
||||
self.tokens = []
|
||||
self._scan()
|
||||
|
||||
def _scan(self):
|
||||
i = 0
|
||||
n = len(self.text)
|
||||
while i < n:
|
||||
m = _TOKEN_RE.match(self.text, i)
|
||||
if not m:
|
||||
# skip whitespace
|
||||
if self.text[i].isspace():
|
||||
i += 1
|
||||
continue
|
||||
raise ValueError(f"Unexpected character {self.text[i]!r} at position {i}")
|
||||
i = m.end()
|
||||
if m.group("num") is not None:
|
||||
self.tokens.append((_NUMBER, m.group("num")))
|
||||
elif m.group("ident") is not None:
|
||||
self.tokens.append((_IDENT, m.group("ident")))
|
||||
elif m.group("op") is not None:
|
||||
op = m.group("op")
|
||||
if op == "(":
|
||||
self.tokens.append((_LPAREN, op))
|
||||
elif op == ")":
|
||||
self.tokens.append((_RPAREN, op))
|
||||
elif op == ",":
|
||||
self.tokens.append((_COMMA, op))
|
||||
else:
|
||||
self.tokens.append((_OP, op))
|
||||
self.tokens.append((_EOF, None))
|
||||
|
||||
def __iter__(self):
|
||||
return iter(self.tokens)
|
||||
|
||||
|
||||
class _Parser:
|
||||
def __init__(self, tokens):
|
||||
self.tokens = tokens
|
||||
self.i = 0
|
||||
|
||||
def peek(self):
|
||||
return self.tokens[self.i]
|
||||
|
||||
def advance(self):
|
||||
t = self.tokens[self.i]
|
||||
self.i += 1
|
||||
return t
|
||||
|
||||
def expect(self, kind, value=None):
|
||||
t = self.peek()
|
||||
if t[0] != kind or (value is not None and t[1] != value):
|
||||
raise ValueError(f"Expected {kind} {value!r} but got {t!r}")
|
||||
return self.advance()
|
||||
|
||||
# Grammar:
|
||||
# expr := term (('+'|'-') term)*
|
||||
# term := power (('*'|'/'|'//'|'%') power)*
|
||||
# power := unary ('**' power)? (right-associative)
|
||||
# unary := ('+'|'-') unary | primary
|
||||
# primary:= NUMBER | IDENT ( '(' args ')' )? | '(' expr ')'
|
||||
# args := expr (',' expr)*
|
||||
|
||||
def parse(self):
|
||||
result = self.expr()
|
||||
if self.peek()[0] != _EOF:
|
||||
raise ValueError(f"Unexpected trailing token {self.peek()!r}")
|
||||
return result
|
||||
|
||||
def expr(self):
|
||||
value = self.term()
|
||||
while self.peek()[0] == _OP and self.peek()[1] in ("+", "-"):
|
||||
op = self.advance()[1]
|
||||
rhs = self.term()
|
||||
if op == "+":
|
||||
value = value + rhs
|
||||
else:
|
||||
value = value - rhs
|
||||
return value
|
||||
|
||||
def term(self):
|
||||
value = self.power()
|
||||
while self.peek()[0] == _OP and self.peek()[1] in ("*", "/", "//", "%"):
|
||||
op = self.advance()[1]
|
||||
rhs = self.power()
|
||||
if op == "*":
|
||||
value = value * rhs
|
||||
elif op == "/":
|
||||
if rhs == 0:
|
||||
raise ValueError("Division by zero")
|
||||
value = value / rhs
|
||||
elif op == "//":
|
||||
if rhs == 0:
|
||||
raise ValueError("Division by zero")
|
||||
value = value // rhs
|
||||
else:
|
||||
if rhs == 0:
|
||||
raise ValueError("Modulo by zero")
|
||||
value = value % rhs
|
||||
return value
|
||||
|
||||
def power(self):
|
||||
base = self.unary()
|
||||
if self.peek()[0] == _OP and self.peek()[1] == "**":
|
||||
self.advance()
|
||||
exp = self.power() # right-assoc
|
||||
return base ** exp
|
||||
return base
|
||||
|
||||
def unary(self):
|
||||
if self.peek()[0] == _OP and self.peek()[1] in ("+", "-"):
|
||||
op = self.advance()[1]
|
||||
val = self.unary()
|
||||
return -val if op == "-" else val
|
||||
return self.primary()
|
||||
|
||||
def primary(self):
|
||||
t = self.peek()
|
||||
if t[0] == _NUMBER:
|
||||
self.advance()
|
||||
s = t[1]
|
||||
return float(s) if ("." in s or "e" in s.lower()) else int(s)
|
||||
if t[0] == _IDENT:
|
||||
name = t[1]
|
||||
self.advance()
|
||||
if self.peek()[0] == _LPAREN:
|
||||
if name not in _FUNCTIONS:
|
||||
raise ValueError(f"Unknown function {name!r}")
|
||||
self.advance() # '('
|
||||
args = [self.expr()]
|
||||
while self.peek()[0] == _COMMA:
|
||||
self.advance()
|
||||
args.append(self.expr())
|
||||
self.expect(_RPAREN)
|
||||
try:
|
||||
return _FUNCTIONS[name](*args)
|
||||
except Exception as exc:
|
||||
raise ValueError(f"Error in {name}: {exc}")
|
||||
if name in _CONSTANTS:
|
||||
return _CONSTANTS[name]
|
||||
raise ValueError(f"Unknown identifier {name!r}")
|
||||
if t[0] == _LPAREN:
|
||||
self.advance()
|
||||
val = self.expr()
|
||||
self.expect(_RPAREN)
|
||||
return val
|
||||
raise ValueError(f"Unexpected token {t!r}")
|
||||
|
||||
|
||||
def evaluate(expr):
|
||||
if not isinstance(expr, str):
|
||||
raise ValueError("Expression must be a string")
|
||||
if not expr.strip():
|
||||
raise ValueError("Empty expression")
|
||||
tokens = _Tokenizer(expr).tokens
|
||||
parser = _Parser(tokens)
|
||||
return parser.parse()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
print(evaluate(" ".join(sys.argv[1:])))
|
||||
@@ -0,0 +1,44 @@
|
||||
import sys
|
||||
sys.path.insert(0, "/root/.hermes/wizard-outputs/beam")
|
||||
from safe_eval import evaluate
|
||||
|
||||
cases = [
|
||||
("2+3*4", 14),
|
||||
("(2+3)*4", 20),
|
||||
("2**3**2", 512), # right assoc
|
||||
("10//3", 3),
|
||||
("10%3", 1),
|
||||
("-5+2", -3),
|
||||
("abs(-7)", 7),
|
||||
("min(3,1,2)", 1),
|
||||
("max(3,1,2)", 3),
|
||||
("round(2.5)", 2), # banker's? python round
|
||||
("sqrt(16)+pi*0", 4.0),
|
||||
("2.5*2", 5.0),
|
||||
("log(e)", 1.0),
|
||||
("sin(0)", 0.0),
|
||||
]
|
||||
for e, exp in cases:
|
||||
got = evaluate(e)
|
||||
assert got == exp, (e, got, exp)
|
||||
print("OK", e, "=", got)
|
||||
|
||||
# error cases
|
||||
errs = ["", "1/0", "import os", "x+1", "__import__('os')", "1 +", "(1+2", "f(1)"]
|
||||
for e in errs:
|
||||
try:
|
||||
evaluate(e)
|
||||
print("FAIL no error:", repr(e))
|
||||
sys.exit(1)
|
||||
except ValueError:
|
||||
print("ERR-OK", repr(e))
|
||||
|
||||
# injection attempts
|
||||
for bad in ["__builtins__", "eval('1')", "().__class__"]:
|
||||
try:
|
||||
evaluate(bad)
|
||||
print("FAIL no error:", repr(bad)); sys.exit(1)
|
||||
except ValueError:
|
||||
print("INJ-OK", repr(bad))
|
||||
|
||||
print("ALL TESTS PASSED")
|
||||
Reference in new issue
Block a user