From 66e6596b5e80ba6067cf2c9888e7b657ffa410ce Mon Sep 17 00:00:00 2001
From: Lin Guo <linsword13@gmail.com>
Date: Thu, 5 Sep 2024 15:51:37 -0700
Subject: [PATCH] Add support for regex match in expander

Can be useful for things like filtering experiments:

```
// run only h3 experiments that use less than 10 nodes
ramble workspace setup --where 're_search("h3.*_node_\d_", {experiment_name})'
```
---
 lib/ramble/docs/workspace_config.rst | 1 +
 lib/ramble/ramble/expander.py        | 7 +++++++
 lib/ramble/ramble/test/expander.py   | 2 ++
 3 files changed, 10 insertions(+)

diff --git a/lib/ramble/docs/workspace_config.rst b/lib/ramble/docs/workspace_config.rst
index d90456a75..066cf976a 100644
--- a/lib/ramble/docs/workspace_config.rst
+++ b/lib/ramble/docs/workspace_config.rst
@@ -174,6 +174,7 @@ Supported functions are:
 * ``simplify_str()`` (convert input string to only alphanumerical characters and dashes)
 * ``randrange`` (from `random.randrange`)
 * ``randint`` (from `random.randint`)
+* ``re_search(regex, str)`` (determine if ``str`` contains pattern ``regex``, based on ``re.search``)
 
 .. _ramble-escaped-variables:
 
diff --git a/lib/ramble/ramble/expander.py b/lib/ramble/ramble/expander.py
index 5051997b1..b0570f455 100644
--- a/lib/ramble/ramble/expander.py
+++ b/lib/ramble/ramble/expander.py
@@ -29,6 +29,12 @@ def _or(a, b):
     return a or b
 
 
+def _re_search(regex, s):
+    import re
+
+    return re.search(regex, s) is not None
+
+
 supported_math_operators = {
     ast.Add: operator.add,
     ast.Sub: operator.sub,
@@ -60,6 +66,7 @@ def _or(a, b):
     "randrange": random.randrange,
     "randint": random.randint,
     "simplify_str": spack.util.naming.simplify_name,
+    "re_search": _re_search,
 }
 
 
diff --git a/lib/ramble/ramble/test/expander.py b/lib/ramble/ramble/test/expander.py
index d344f3c08..7fca7f34b 100644
--- a/lib/ramble/ramble/test/expander.py
+++ b/lib/ramble/ramble/test/expander.py
@@ -73,6 +73,8 @@ def exp_dict():
         ('"2.1.1" in ["2.1.1", "3.1.1", "4.2.1"]', "True", set(), 1),
         ('"2.1.2" in ["2.1.1", "3.1.1", "4.2.1"]', "False", set(), 1),
         ("{test_mask}", "0x0", {"test_mask"}, 1),
+        ('re_search(r"bz$", {experiment_name})', "False", set(), 1),
+        ('re_search(r"_fo+", {env_name})', "True", set(), 1),
     ],
 )
 def test_expansions(input, output, no_expand_vars, passes):