Merge pull request #140 from Deltares/prototype_REACT

Prototype react: make numpy and math available in formula based rule

Author: HiddeElzinga

decoimpact/business/entities/rules/formula_rule.py

@@ -11,9 +11,12 @@
     Formula Rule
 """
 
+# Import safe modules
+import math
 from argparse import ArgumentError as _ArgumentError
 from typing import Dict, List
 
+import numpy
 from RestrictedPython import compile_restricted as _compile_restricted
 from RestrictedPython import safe_builtins as _safe_builtins
 
@@ -85,23 +88,22 @@ def execute(self, values: Dict[str, float], logger: ILogger) -> float:
         return float(local_variables[self.formula_output_name])
 
     def _setup_environment(self):
-        self._safe_modules = frozenset(
-            (
-                "math",
-                "numpy",
-            )
-        )
+        # use standard libraries that are considered safe
+        self._safe_modules_dict = {
+            "math": math,
+            "numpy": numpy,
+        }
 
         # Global data available in restricted code
-        self._global_variables = (
-            {  # MDK: THIS NEEDS TO CHANGE TO A MORE GENERAL APPROACH
-                "__builtins__": {**_safe_builtins, "__import__": self._safe_import},
-            }
-        )
+        self._global_variables = {
+            "__builtins__": {**_safe_builtins, "__import__": self._safe_import},
+            **self._safe_modules_dict,
+        }
+
         self._byte_code = None
 
     def _safe_import(self, name, *args, **kwargs):
         # Redefine import, to only import from safe modules
-        if name not in self._safe_modules:
+        if name not in self._safe_modules_dict:
             raise _ArgumentError(None, f"Importing {name!r} is not allowed!")
         return __import__(name, *args, **kwargs)

docs/manual/rules/formula_rule.md

@@ -69,3 +69,5 @@ When a formula results in a boolean, it will be converted to a float result. Mea
 
 For more information on these operators click [here](https://www.w3schools.com/python/python_operators.asp).
 
+It is also possible to use functions of the libraries math and numpy. These are accessible by calling their full module names inside the formula, for instance:
+"numpy.where(water_depth > 1.0)" or "math.ceil(water_level)".

tests/business/entities/rules/test_formula_rule.py

@@ -67,6 +67,36 @@ def test_execute_multiplying_value_arrays():
     assert math.isclose(result_value, 6.0, abs_tol=1e-9)
 
 
+def test_execute_numpy_value_arrays():
+    """Test formula on value_arrays of a RuleBase"""
+
+    # Arrange
+    logger = Mock(ILogger)
+    rule = FormulaRule("test", ["val1"], "val1 * numpy.size(numpy.array([1, 3]))")
+    values = {"val1": 1.0}
+
+    # Act
+    result_value = rule.execute(values, logger)
+
+    # Assert
+    assert math.isclose(result_value, 2.0, abs_tol=1e-9)
+
+
+def test_execute_math_value_arrays():
+    """Test formula on value_arrays of a RuleBase"""
+
+    # Arrange
+    logger = Mock(ILogger)
+    rule = FormulaRule("test", ["val1"], "val1 * math.isqrt(9)")
+    values = {"val1": 2.0}
+
+    # Act
+    result_value = rule.execute(values, logger)
+
+    # Assert
+    assert math.isclose(result_value, 6.0, abs_tol=1e-9)
+
+
 @pytest.mark.parametrize(
     "input_value1, input_value2, expected_output_value",
     [(0.5, 10, 0.0), (11, 1.5, 1.0)],