troupe-standup/wizard-work/beam/safe_eval.py

"""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:])))