91 lines
2.3 KiB
Python
91 lines
2.3 KiB
Python
import ast
|
|
import datetime
|
|
import math
|
|
import random
|
|
from typing import Any, Dict, Set
|
|
|
|
|
|
SAFE_BUILTINS = {
|
|
"abs": abs,
|
|
"all": all,
|
|
"any": any,
|
|
"bool": bool,
|
|
"dict": dict,
|
|
"divmod": divmod,
|
|
"enumerate": enumerate,
|
|
"filter": filter,
|
|
"float": float,
|
|
"format": format,
|
|
"int": int,
|
|
"isinstance": isinstance,
|
|
"len": len,
|
|
"list": list,
|
|
"map": map,
|
|
"max": max,
|
|
"min": min,
|
|
"pow": pow,
|
|
"range": range,
|
|
"round": round,
|
|
"set": set,
|
|
"sorted": sorted,
|
|
"str": str,
|
|
"sum": sum,
|
|
"tuple": tuple,
|
|
"zip": zip,
|
|
}
|
|
|
|
SAFE_GLOBALS: Dict[str, Any] = {
|
|
"__builtins__": SAFE_BUILTINS,
|
|
"math": math,
|
|
"random": random,
|
|
"datetime": datetime.datetime,
|
|
"date": datetime.date,
|
|
"timedelta": datetime.timedelta,
|
|
}
|
|
|
|
|
|
class FormulaSyntaxError(Exception):
|
|
pass
|
|
|
|
|
|
class FormulaRuntimeError(Exception):
|
|
pass
|
|
|
|
|
|
def extract_variable_names(expression_str: str) -> Set[str]:
|
|
if not expression_str or not expression_str.strip():
|
|
return set()
|
|
try:
|
|
tree = ast.parse(expression_str, mode="eval")
|
|
except SyntaxError as e:
|
|
raise FormulaSyntaxError(f"Syntax error in expression '{expression_str}': {e.msg}") from e
|
|
|
|
reserved = set(SAFE_GLOBALS.keys()) | set(SAFE_BUILTINS.keys())
|
|
vars_found: Set[str] = set()
|
|
for node in ast.walk(tree):
|
|
if isinstance(node, ast.Name):
|
|
if node.id not in reserved:
|
|
vars_found.add(node.id)
|
|
return vars_found
|
|
|
|
|
|
class CompiledExpression:
|
|
def __init__(self, expression_str: str, target_name: str = "<expr>"):
|
|
self.raw_expr = expression_str.strip()
|
|
self.target_name = target_name
|
|
self.referenced_vars = extract_variable_names(self.raw_expr)
|
|
try:
|
|
self.code = compile(self.raw_expr, f"<{target_name}>", "eval")
|
|
except SyntaxError as e:
|
|
raise FormulaSyntaxError(
|
|
f"Syntax error in '{target_name}': {e.msg} at line {e.lineno}, offset {e.offset}"
|
|
) from e
|
|
|
|
def evaluate(self, context: Dict[str, Any]) -> Any:
|
|
try:
|
|
return eval(self.code, SAFE_GLOBALS, context)
|
|
except Exception as e:
|
|
raise FormulaRuntimeError(
|
|
f"Evaluation error in '{self.target_name}' using context {context}: {type(e).__name__} - {e}"
|
|
) from e
|