diff --git a/monai/utils/safeeval.py b/monai/utils/safeeval.py index dd357601a8..e8f0bcb1ae 100644 --- a/monai/utils/safeeval.py +++ b/monai/utils/safeeval.py @@ -47,11 +47,12 @@ def __init__(self, int_type_str: str, float_type_str: str): self.float_type_str = float_type_str def visit_Constant(self, node): - if isinstance(node.value, (int, float)): - type_str = self.int_type_str if isinstance(node.value, int) else self.float_type_str - return ast.parse(f"{type_str}({node.value})") - - return node + if isinstance(node.value, bool) or not isinstance(node.value, (int, float)): + return node + type_str = self.int_type_str if isinstance(node.value, int) else self.float_type_str + func_node = ast.parse(type_str, mode="eval").body + call_node = ast.Call(func=func_node, args=[ast.Constant(value=node.value)], keywords=[]) + return ast.copy_location(call_node, node) def safe_eval( @@ -74,7 +75,7 @@ def safe_eval( by `int_type_str` and `float_type_str`. These are expected to be constructor names prefixed with `np.` as Numpy will be present in the expression global variables under that name. The values can be changed to other types if needed, such as "int64". One advantage of doing this is to avoid denial-of-service attacks by attempting to evaluate - an expressoini which is incredibly slow under native Python but fast (though potentially erroneous) under Numpy. + an expression which is incredibly slow under native Python but fast (though potentially erroneous) under Numpy. Args: expr: expression to evaluate, this will be stripped before parsing to avoid indentation complaints @@ -101,6 +102,7 @@ def safe_eval( if rewrite_np: parsed = _RewriteConstNp(int_type_str, float_type_str).visit(parsed) - locals_vars = {"np": np, **(locals_vars or {})} + ast.fix_missing_locations(parsed) + locals_vars = {**(locals_vars or {}), "np": np} - return eval(expr, dict(globals_vars) if globals_vars else None, locals_vars) + return eval(compile(parsed, "", "eval"), dict(globals_vars) if globals_vars else None, locals_vars) diff --git a/tests/utils/test_safe_eval.py b/tests/utils/test_safe_eval.py index 836dcb90b2..b0fd5d8ba6 100644 --- a/tests/utils/test_safe_eval.py +++ b/tests/utils/test_safe_eval.py @@ -14,6 +14,7 @@ import ast import unittest +import numpy as np from parameterized import parameterized from monai.utils import safe_eval @@ -62,6 +63,34 @@ def test_allowed_types(self): with self.assertRaises(ValueError): safe_eval("1*2", allowed_types=allowed) + def test_rewrite_np_produces_numpy_types(self): + """Test that rewrite_np wraps literals in numpy types.""" + result = safe_eval("2 + 3", rewrite_np=True) + self.assertIsInstance(result, np.integer) + + result = safe_eval("2.5 + 1.5", rewrite_np=True) + self.assertIsInstance(result, np.floating) + + def test_rewrite_np_large_exponent(self): + """Test that rewrite_np prevents slow native-Python exponentiation.""" + # Under native Python, 9**9**9 produces a ~369-million-digit integer; + # under np.int32 it overflows and completes almost instantly. + result = safe_eval("9**9**9", rewrite_np=True) + self.assertIsInstance(result, np.integer) + + def test_rewrite_np_preserves_bool(self): + """Test that rewrite_np does not wrap bool constants.""" + result = safe_eval("True", rewrite_np=True) + self.assertIs(result, True) + + result = safe_eval("False", rewrite_np=True) + self.assertIs(result, False) + + def test_rewrite_np_inf_constant(self): + """Test that rewrite_np handles inf/nan literals correctly.""" + result = safe_eval("1e309", rewrite_np=True) + self.assertIsInstance(result, np.floating) + if __name__ == "__main__": unittest.main()