Merge branch 'remember' into 'main'

Remember data for use in later steps

Closes #107

See merge request larswirzenius/subplot!103

2 files changed, 69 insertions, 0 deletions

diff --git a/templates/python/context.py b/templates/python/context.py
index 9ecae19..d61316e 100644
--- a/templates/python/context.py
+++ b/templates/python/context.py
@@ -1,4 +1,5 @@
 import logging
+import re
 
 
 # Store context between steps.
@@ -35,6 +36,33 @@ class Context:
             logging.debug(f"Context: declared {name}")
         return self._ns[name]
 
+    def remember_value(self, name, value):
+        ns = self.declare("_values")
+        if name in ns:
+            raise KeyError(name)
+        ns[name] = value
+
+    def recall_value(self, name):
+        ns = self.declare("_values")
+        if name not in ns:
+            raise KeyError(name)
+        return ns[name]
+
+    def expand_values(self, pattern):
+        parts = []
+        while pattern:
+            m = re.search(r"(?<!\$)\$\{(?P<name>\S*)\}", pattern)
+            if not m:
+                parts.append(pattern)
+                break
+            name = m.group("name")
+            if not name:
+                raise KeyError("empty name in expansion")
+            value = self.recall_value(name)
+            parts.append(value)
+            pattern = pattern[m.end() :]
+        return "".join(parts)
+
 
 class NameSpace:
     def __init__(self, name):
diff --git a/templates/python/context_tests.py b/templates/python/context_tests.py
index d878084..c91350e 100644
--- a/templates/python/context_tests.py
+++ b/templates/python/context_tests.py
@@ -52,6 +52,47 @@ class ContextTests(unittest.TestCase):
         self.assertTrue("bar" in repr(ctx))
 
 
+class ContextMemoryTests(unittest.TestCase):
+    def test_recall_raises_exception_for_unremembered_value(self):
+        ctx = Context()
+        with self.assertRaises(KeyError):
+            ctx.recall_value("foo")
+
+    def test_recall_returns_remembered_value(self):
+        ctx = Context()
+        ctx.remember_value("foo", "bar")
+        self.assertEqual(ctx.recall_value("foo"), "bar")
+
+    def test_remember_raises_exception_for_previously_remembered(self):
+        ctx = Context()
+        ctx.remember_value("foo", "bar")
+        with self.assertRaises(KeyError):
+            ctx.remember_value("foo", "bar")
+
+    def test_expand_returns_pattern_without_values_as_is(self):
+        ctx = Context()
+        self.assertEqual(ctx.expand_values("foo"), "foo")
+
+    def test_expand_allows_double_dollar_escapes(self):
+        ctx = Context()
+        self.assertEqual(ctx.expand_values("$${foo}"), "$${foo}")
+
+    def test_expand_raises_exception_for_empty_name_expansion_as_is(self):
+        ctx = Context()
+        with self.assertRaises(KeyError):
+            ctx.expand_values("${}")
+
+    def test_expand_raises_error_for_unrememebered_values(self):
+        ctx = Context()
+        with self.assertRaises(KeyError):
+            ctx.expand_values("${foo}")
+
+    def test_expands_rememebered_values(self):
+        ctx = Context()
+        ctx.remember_value("foo", "bar")
+        self.assertEqual(ctx.expand_values("${foo}"), "bar")
+
+
 class ContextNamepaceTests(unittest.TestCase):
     def test_explicit_namespaces_are_empty_dicts_initially(self):
         ctx = Context()