Add arithmetic expression support in batch matrix parameters

qaboard/iterators.py

@@ -3,9 +3,11 @@
 """
 import os
 import re
+import ast
 import sys
 import json
 import numbers
+import operator
 import fnmatch
 import traceback
 from copy import deepcopy
@@ -308,6 +310,96 @@ def deep_interpolate(value, replaced: str, to_value):
     return value
 
 
+# Arithmetic expressions like `${{ 168 * matrix.gain }}`, inspired by GitHub Actions
+expression_regex = re.compile(r'\$\{\{(.*?)\}\}')
+
+safe_eval_binops = {
+  ast.Add: operator.add,
+  ast.Sub: operator.sub,
+  ast.Mult: operator.mul,
+  ast.Div: operator.truediv,
+  ast.FloorDiv: operator.floordiv,
+  ast.Mod: operator.mod,
+  ast.Pow: operator.pow,
+}
+safe_eval_unaryops = {ast.UAdd: operator.pos, ast.USub: operator.neg}
+safe_eval_functions = {'abs': abs, 'int': int, 'float': float, 'round': round, 'min': min, 'max': max}
+
+def safe_eval(expression: str, variables: Dict):
+  """Evaluates arithmetic expressions like "168 * matrix.gain". Anything except literals,
+  variable lookups, arithmetic operators and abs/int/float/round/min/max raises ValueError."""
+  def eval_node(node):
+    node_type = type(node).__name__
+    if isinstance(node, ast.Expression):
+      return eval_node(node.body)
+    if isinstance(node, ast.Constant):
+      return node.value
+    if node_type == 'Num': # python3.7
+      return node.n
+    if node_type == 'Str': # python3.7
+      return node.s
+    if isinstance(node, ast.Name):
+      if node.id in variables:
+        return variables[node.id]
+      if node.id in safe_eval_functions:
+        return safe_eval_functions[node.id]
+      raise ValueError(f'unknown name "{node.id}"')
+    if isinstance(node, ast.Attribute):
+      obj = eval_node(node.value)
+      if isinstance(obj, dict) and node.attr in obj:
+        return obj[node.attr]
+      raise ValueError(f'unknown attribute "{node.attr}"')
+    if isinstance(node, ast.Subscript):
+      obj = eval_node(node.value)
+      key = node.slice
+      if type(key).__name__ == 'Index': # python<3.9
+        key = key.value
+      return obj[eval_node(key)]
+    if isinstance(node, ast.BinOp) and type(node.op) in safe_eval_binops:
+      return safe_eval_binops[type(node.op)](eval_node(node.left), eval_node(node.right))
+    if isinstance(node, ast.UnaryOp) and type(node.op) in safe_eval_unaryops:
+      return safe_eval_unaryops[type(node.op)](eval_node(node.operand))
+    if isinstance(node, ast.Call):
+      func = eval_node(node.func)
+      if func not in safe_eval_functions.values():
+        raise ValueError(f'only calls to {"/".join(safe_eval_functions)} are allowed')
+      return func(*[eval_node(a) for a in node.args], **{kw.arg: eval_node(kw.value) for kw in node.keywords})
+    raise ValueError(f'unsupported syntax ({node_type})')
+  return eval_node(ast.parse(expression.strip(), mode='eval'))
+
+
+def evaluate_expressions(value, variables: Dict):
+  """Evaluates recursively the `${{ <expression> }}` blocks found in strings.
+  If a string is exactly one expression, it keeps the result's type (e.g. int),
+  otherwise results are formatted back into the string (floats with integral
+  values are formatted as ints, so `${{ 168 * 1.5 }}` gives "252" not "252.0").
+  Unlike the lenient `${matrix.param}` interpolation, errors here raise ValueError:
+  expressions are always intentional, so failures should never pass silently."""
+  if isinstance(value, dict):
+    return {k: evaluate_expressions(v, variables) for k, v in value.items()}
+  if isinstance(value, list):
+    return [evaluate_expressions(v, variables) for v in value]
+  if not isinstance(value, str):
+    return value
+  def evaluate(match):
+    expression = match.group(1)
+    try:
+      return safe_eval(expression, variables)
+    except Exception as e:
+      message = f'ERROR: Cannot evaluate the expression <${{{{{expression}}}}}>: {e}'
+      click.secho(message, fg='red', bold=True, err=True)
+      raise ValueError(message) from e
+  full_match = expression_regex.fullmatch(value.strip())
+  if full_match:
+    return evaluate(full_match)
+  def evaluate_str(match):
+    result = evaluate(match)
+    if isinstance(result, float) and result.is_integer():
+      return str(int(result))
+    return str(result)
+  return expression_regex.sub(evaluate_str, value)
+
+
 def iter_batch(batch: Dict, default_run_context: RunContext, qatools_config, default_inputs_settings, debug):
     # Happens often when there is an orphan "my-batch:" in in the yaml file
     if batch is None:
@@ -356,10 +448,14 @@ def iter_batch(batch: Dict, default_run_context: RunContext, qatools_config, def
             matrix_run_context.configurations.append(matrix_config)
           else:
             matrix_run_context.configurations = matrix_config
+        matrix_params = {}
         for param, value in matrix.items():
           if param in ['configuration', 'configurations', 'configs', 'platform']:
             continue
+          matrix_params[param] = value
           matrix_run_context.configurations = deep_interpolate(matrix_run_context.configurations, 'matrix', {param: value})
+        # arithmetic expressions are evaluated after all matrix parameters are known
+        matrix_run_context.configurations = evaluate_expressions(matrix_run_context.configurations, {'matrix': matrix_params})
         yield from iter_batch(batch_, matrix_run_context, qatools_config, default_inputs_settings, debug)
       return
 

tests/test_iterators.py

@@ -82,6 +82,32 @@ def test_interpolation(self):
       "object abcdef"
     )
 
+  def test_expressions(self):
+    from qaboard.iterators import evaluate_expressions
+    variables = {'matrix': {'gain': 2}}
+    # full-string expressions keep their type
+    self.assertEqual(evaluate_expressions("${{ 168 * matrix.gain }}", variables), 336)
+    self.assertEqual(evaluate_expressions("${{ matrix.gain / 4 }}", variables), 0.5)
+    self.assertEqual(evaluate_expressions("${{ min(2 ** matrix.gain, 3) }}", variables), 3)
+    self.assertEqual(evaluate_expressions("${{ matrix['gain'] + 1 }}", variables), 3)
+    # expressions embedded in strings are formatted back, recursively in dicts/lists
+    self.assertEqual(
+      evaluate_expressions({"SMAP.bp_th": "[${{ 168 * matrix.gain }}, 200, ${{ matrix.gain }}]"}, variables),
+      {"SMAP.bp_th": "[336, 200, 2]"}
+    )
+    # integral floats are formatted as ints
+    self.assertEqual(evaluate_expressions("[${{ 168 * matrix.gain }}]", {'matrix': {'gain': 1.5}}), "[252]")
+    # strings without expressions are left alone
+    self.assertEqual(evaluate_expressions("plain ${matrix.gain}", variables), "plain ${matrix.gain}")
+    self.assertEqual(evaluate_expressions(42, variables), 42)
+    # errors are loud, not silent
+    with self.assertRaises(ValueError):
+      evaluate_expressions("${{ matrix.typo }}", variables)
+    with self.assertRaises(ValueError):
+      evaluate_expressions("${{ __import__('os').system('true') }}", variables)
+    with self.assertRaises(ValueError):
+      evaluate_expressions("${{ 1 if 1 else 2 }}", variables)
+
   def test_match(self):
     from qaboard.iterators import match
     metadata = {"Sensor": "HM4"}
@@ -208,6 +234,11 @@ def get_batch(batch):
     batches = get_batch('matrix-interpolate-2')
     self.assertEqual(len(batches), 4)
 
+    batches = get_batch('matrix-expressions')
+    self.assertEqual(len(batches), 2)
+    self.assertEqual(batches[0].configurations, ['base', {"SMAP.bp_th": "[168, 200]"}, {"threshold": 10}])
+    self.assertEqual(batches[1].configurations, ['base', {"SMAP.bp_th": "[336, 200]"}, {"threshold": 20}])
+
 
 
 sample_batches_yaml = """
@@ -360,6 +391,16 @@ def get_batch(batch):
     - config-v${matrix.version}
     - version: v${matrix.version}
 
+matrix-expressions:
+  inputs:
+  - a.txt
+  matrix:
+    gain: [1, 2]
+  configurations:
+    - base
+    - SMAP.bp_th: "[${{ 168 * matrix.gain }}, 200]"
+    - threshold: ${{ 10 * matrix.gain }}
+
 matrix-interpolate-2:
   inputs:
   - a.txt

website/docs/batches-running-on-multiple-inputs.md

@@ -193,6 +193,27 @@ my-batch-matrix-object:
     - flags: "-x=${matrix.point[x]} -y=${matrix.point[y]}"
 ```
 
+### Arithmetic expressions
+Sometimes parameters are correlated: a threshold that scales with the gain, a size that depends on a ratio... Instead of maintaining parallel lists of pre-computed values, you can write arithmetic expressions with the `${{ <expression> }}` syntax, referencing matrix parameters as `matrix.variable`:
+
+```yaml
+my-batch-derived-parameters:
+  inputs:
+  - image.raw
+  matrix:
+    gain: [1, 2, 4]
+  configs:
+    - workspace/config_gain${matrix.gain}.cde
+    - SMAP.bp_th_red_min_curve_exp_0: "[${{ 168 * matrix.gain }}, 200, 240, 280, 360, 600, 600]"
+    - SMAP.other_threshold: ${{ min(50 * matrix.gain, 100) }}
+
+# => will run with SMAP.bp_th_red_min_curve_exp_0 == "[168, ...]", "[336, ...]" and "[672, ...]"
+```
+
+Expressions support the usual arithmetic operators (`+ - * / // % **`), parentheses, and calls to `abs/int/float/round/min/max` — nothing else, by design. If an expression is the whole value, the result keeps its type (*e.g.* integer); when embedded in a longer string it is formatted back into it (floats with integral values like `252.0` are written as `252`).
+
+> Unlike `${matrix.variable}` interpolation, which leaves unknown placeholders untouched, a `${{ }}` expression that cannot be evaluated stops the batch with an error: expressions are always intentional, so typos fail loudly instead of silently sending a literal string to your runs.
+
 ## Aliases for groups of batches
 For convenience you can define aliases for batches you often run together. For instance you can do: