Repository: radiasoft/sirepo

Issue:
Add array index evaluation to CodeVar
As a counterpart to #6503, handle expressions of the form <array variable>[<index>], and modify the existing test case.



Discussion:

Patch:
diff --git a/sirepo/template/code_variable.py b/sirepo/template/code_variable.py
index 7eccf03c6..49b8b4569 100644
--- a/sirepo/template/code_variable.py
+++ b/sirepo/template/code_variable.py
@@ -173,6 +173,8 @@ class CodeVar:
                 expr = re.sub(r"\^", "**", expr)
                 rpn = cls.__parse_expr_infix(expr)
                 expr = rpn
+            else:
+                expr = float(expr)
         except Exception as e:
             pass
         return expr
@@ -343,15 +345,19 @@ class PurePythonEval:
     def eval_var(self, expr, depends, variables):
         variables = variables.copy()
         for d in depends:
-            v, err = PurePythonEval.__eval_python_stack(
+            # recurse eval_var, but with empty dependencies
+            v, err = PurePythonEval.eval_var(
                 self,
-                variables[d],
+                self.__eval_indexed_variable(variables[d], variables),
+                {},
                 variables,
             )
             if err:
                 return None, err
             variables[d] = v
-        return self.__eval_python_stack(expr, variables)
+        return self.__eval_python_stack(
+            self.__eval_indexed_variable(expr, variables), variables
+        )
 
     @classmethod
     def postfix_to_infix(cls, expr):
@@ -387,6 +393,21 @@ class PurePythonEval:
                 stack.append(v)
         return __strip_parens(stack[-1])
 
+    def __eval_indexed_variable(self, expr, variables):
+        if isinstance(expr, list):
+            return [self.__eval_indexed_variable(e, variables) for e in expr]
+        r = rf"(.*)({'|'.join(list(variables.keys()))})\s*\[\s*(\d+)\s*\]"
+        if not re.match(r, str(expr)):
+            return CodeVar.infix_to_postfix(expr)
+        return self.__eval_indexed_variable(
+            re.sub(
+                r,
+                lambda m: m.group(1) + str(variables[m.group(2)][int(m.group(3))]),
+                expr,
+            ),
+            variables,
+        )
+
     def __eval_python_stack(self, expr, variables):
         if not CodeVar.is_var_value(expr):
             return expr, None


Test patch:
diff --git a/tests/template/code_variable_test.py b/tests/template/code_variable_test.py
index f97c638c1..c3d9ba5b1 100644
--- a/tests/template/code_variable_test.py
+++ b/tests/template/code_variable_test.py
@@ -47,7 +47,7 @@ def test_arrays():
     pkeq([124, 125, 126], code_var.eval_var("expressionArray")[0])
     pkeq([[1, 2, 3], [4, 5], [129]], code_var.eval_var("nestedArray")[0])
     # expressions containing arrays are not currently handled
-    pkeq("unknown token: numberArray[0]", code_var.eval_var("arrayExpression")[1])
+    pkeq(2, code_var.eval_var("arrayExpression")[0])
     pkeq("unknown token: y", code_var.eval_var("invalidExpressionArray")[1])