test_safe_eval.py

# Copyright (c) MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import ast
import unittest

from parameterized import parameterized

from monai.utils import safe_eval

GOOD_EXPRS = [
    ("1+2", None, None, 3),
    ("     1    +     2   ", None, None, 3),
    ("1+2+x", {"x": 4}, None, 7),
    ("1+2+x", None, {"x": 4}, 7),
    ("1*2+x", {"x": 4}, None, 6),
    ("(1+2)*3", None, None, 9),
    ("foo+bar", {"foo": 1030}, {"bar": 204}, 1234),
]

BAD_EXPRS = [("foo()",), ("foo.bar",), ("foo[123]",), ("(1,2)",), ("[3,4]",), ("int.__class__.__init__.__globals__",)]


class TestSafeEval(unittest.TestCase):
    @parameterized.expand(GOOD_EXPRS)
    def test_good_exprs(self, expr, globals_vars, locals_vars, expected):
        """Test valid expressions with globals/locals evaluate to correct values."""
        result = safe_eval(expr, globals_vars, locals_vars)
        self.assertEqual(result, expected)

    @parameterized.expand(BAD_EXPRS)
    def test_bad_exprs(self, expr):
        """Test bad expressions correctly raise ValueError."""
        with self.assertRaises(ValueError):
            safe_eval(expr)

    def test_allowed_types(self):
        """Test restricting the allowed list of types."""
        allowed = [ast.Expression, ast.Constant, ast.BinOp, ast.Add]
        result = safe_eval("1+2", allowed_types=allowed)
        self.assertEqual(result, 3)

        with self.assertRaises(ValueError):
            safe_eval("1*2", allowed_types=allowed)


if __name__ == "__main__":
    unittest.main()