Fix bracket assignment

README.md

@@ -8,6 +8,7 @@ This is a fork of the original [`airspeed`](https://github.com/purcell/airspeed)
 ⚠️ Note: This fork of `airspeed` focuses on providing maximum parity with AWS' implementation of Velocity templates (used in, e.g., API Gateway or AppSync). In some cases, the behavior may diverge from the VTL spec, or from the Velocity [reference implementation](https://velocity.apache.org/download.cgi).
 
 ## Change Log:
+* v0.6.8: Added support for bracket in Assignment with `#set`; Added support for `Array.set`
 * v0.6.7: fix support for floating point starting with a decimal; Implement `REPLACE_FORMAL_TEXT` to allow bypassing silent behavior of `FormalReference` element.
 * v0.6.6: add support for `$string.matches( $pattern )`; fix bug where some escaped character would prevent string matching
 * v0.6.5: handle `$map.put('key', null)` correctly

airspeed/operators.py

@@ -32,6 +32,16 @@ def dict_put(self, key, value):
     return existing
 
 
+def array_set(self, index, value):
+    try:
+        existing = self[index]
+    except IndexError:
+        raise TemplateExecutionError
+
+    self[index] = value
+    return existing
+
+
 def dict_to_string(obj: dict) -> str:
     return "{" + ", ".join([f"{k}={v}" for k, v in obj.items()]) + "}"
 
@@ -52,6 +62,7 @@ def dict_to_string(obj: dict) -> str:
         "contains": lambda self, value: value in self,
         "add": lambda self, value: self.append(value),
         "isEmpty": lambda self: len(self) == 0,
+        "set": array_set,
     },
     dict: {
         "put": dict_put,
@@ -1068,14 +1079,47 @@ def evaluate_raw(self, stream, namespace, loader):
 # set($one.two().three = something)
 # yet
 class Assignment(_Element):
+    terms: list
     START = re.compile(
-        r"\s*\(\s*\$(\w*(?:\.[\w-]+|\[\"\$\w+\"\]*)*)\s*=\s*(.*)$", re.S + re.I
+        # first group stops at the first '$' encountered. self.end will be set at the first char of the variable
+        # Currently supported in assignment are: `$root`, `.dot`. `["bracket"]`, `[$var]` and `["$quoted_var"]`
+        r"\s*(\(\s*\$)(\w*(?:\.[\w-]+|\[\"\$?\w+\"]|\[\$\w+])*\s*=\s*.*)$",
+        re.S + re.I,
     )
     END = re.compile(r"\s*\)(?:[ \t]*\r?\n)?(.*)$", re.S + re.M)
+    # Allows us to match all supported terms. We are also matching on `=` so we can exit
+    TERMS = re.compile(r"(\.?\w+|\[[\"$]*\w+\"?]|=)", re.S + re.I)
+    TERMS_END = re.compile(r"\s*=\s*(.*)$", re.S)
 
     def parse(self):
-        (var_name,) = self.identity_match(self.START)
-        self.terms = var_name.split(".")
+        self.identity_match(self.START)
+
+        self.terms = []
+        for term_match in self.TERMS.finditer(self._full_text, self.start):
+            term = term_match.group(0)
+            if term == "=":
+                # If we matched the `=` we have gone through the whole variable definition
+                break
+            if term.startswith("."):
+                # handles .dot
+                self.end += len(term)
+                self.terms.append(term[1:])
+            elif "$" in term:
+                # handles ["$quoted_var"] and [$var]
+                # skipping over '['
+                self.end += 1
+                # `Value` handles a lot more than we need, but since we are pretty restrictive on the
+                # `identity_match`, it shouldn't be an issue. If it comes up as a problem in the future we can
+                # restrict the list further
+                self.terms.append(self.require_next_element(Value, "value"))
+                # skipping over ']'
+                self.end += 1
+            else:
+                # handles ["bracket"] and root
+                self.end += len(term)
+                self.terms.append(term.strip('[]"'))
+
+        self.require_match(self.TERMS_END, "=")
         self.value = self.require_next_element(Expression, "expression")
         self.require_match(self.END, ")")
 
@@ -1086,8 +1130,14 @@ def evaluate_raw(self, stream, namespace, loader):
         else:
             cur = namespace
             for term in self.terms[:-1]:
-                cur = cur[term]
-            cur[self.terms[-1]] = val
+                cur = cur[self._calculate_term(term, namespace, loader)]
+            cur[self._calculate_term(self.terms[-1], namespace, loader)] = val
+
+    @staticmethod
+    def _calculate_term(term, namespace, loader):
+        if isinstance(term, Value):
+            return term.calculate(namespace, loader)
+        return term
 
 
 class EvaluateDirective(_Element):

setup.py

@@ -5,7 +5,7 @@
 
 setup(
     name="airspeed-ext",
-    version="0.6.7",
+    version="0.6.8",
     description=(
         "Airspeed is a powerful and easy-to-use templating engine "
         "for Python that aims for a high level of compatibility "

tests/test_templating.py

@@ -808,6 +808,63 @@ def test_dict_put_item(self, test_render):
         )
         test_render(template, {"test_dict": {"k": "initial value"}})
 
+    def test_dict_assign_item_via_brackets(self, test_render):
+        # The specified element is set with the given value.
+        # Velocity tries first the 'set' method on the element, then 'put' to make the assignment.
+        template = (
+            "#set($test_dict = {} )"
+            "#set($key = 'bar')"
+            "#set( $test_dict[$key] = 'foo' )"
+            "$test_dict.toString()"
+        )
+        test_render(template, {})
+
+    def test_dict_assign_quoted_item_via_brackets(self, test_render):
+        # The specified element is set with the given value.
+        # Velocity tries first the 'set' method on the element, then 'put' to make the assignment.
+        template = (
+            "#set($test_dict = {} )"
+            "#set($key = 'bar')"
+            "#set( $test_dict[\"$key\"] = 'foo' )"
+            "$test_dict.toString()"
+        )
+        test_render(template, {})
+
+    def test_dict_assign_nested_items_via_brackets(self, test_render):
+        template = (
+            "#set($test_dict = {} )"
+            "#set($key = 'bar')"
+            '#set( $test_dict["$key"] = {} )'
+            '#set($test_dict["$key"].nested = "foo")'
+            '$test_dict.bar["nested"]'
+        )
+        test_render(template, {})
+
+    def test_dict_assign_var_not_found(self, test_render):
+        template = (
+            "#set($test_dict = {})"
+            '#set($test_dict["$key"] = "foo")'
+            "$test_dict.toString()"
+        )
+        test_render(template, {})
+
+    def test_array_assign_item_via_brackets(self, test_render):
+        template = (
+            "#set($test_array = ['one', 'two', 'three'] )"
+            "#set($i = 1)"
+            "#set( $test_array[$i] = 'foo' )"
+            "#foreach ($item in $test_array)$item#end"
+        )
+        test_render(template, {})
+
+    def test_array_set_item(self, test_render):
+        template = (
+            "#set($test_array = ['one', 'two', 'three'] )"
+            "$test_array.set(1, 'foo')"
+            "#foreach ($item in $test_array)$item#end"
+        )
+        test_render(template, {})
+
     def test_put_null_in_map(self, test_render):
         template = r"""
             #set( $myMap = {} )
@@ -1494,6 +1551,15 @@ def test_get_position_strings_in_syntax_error_when_newline_before_error(
         else:
             pytest.fail("expected error")
 
+    def test_array_set_item_outside_range(self, test_render):
+        template = airspeed.Template(
+            "#set($test_array = ['one', 'two', 'three'] )"
+            "$test_array.set(5, 'foo')"
+            "$test_array"
+        )
+        with pytest.raises(airspeed.TemplateExecutionError):
+            template.merge({})
+
 
 @pytest.mark.skip(reason="Invalid syntax, failing against VTL CLI and/or AWS")
 class TestInvalidCases:

tests/test_templating.snapshot.json

@@ -1097,7 +1097,7 @@
     }
   },
   "tests/test_templating.py::TestTemplating::test_string_matches_full_date": {
-    "recorded-date": "24-10-2024, 21:58:12",
+    "recorded-date": "16-05-2025, 18:29:38",
     "recorded-content": {
       "render-result-1-cli": "true",
       "render-result-1": "true",
@@ -1111,5 +1111,47 @@
       "render-result-1-cli": "0.05",
       "render-result-1": "0.05"
     }
+  },
+  "tests/test_templating.py::TestTemplating::test_dict_assign_item_via_brackets": {
+    "recorded-date": "28-05-2025, 11:06:27",
+    "recorded-content": {
+      "render-result-1-cli": "{bar=foo}",
+      "render-result-1": "{bar=foo}"
+    }
+  },
+  "tests/test_templating.py::TestTemplating::test_array_assign_item_via_brackets": {
+    "recorded-date": "16-05-2025, 19:14:01",
+    "recorded-content": {
+      "render-result-1-cli": "onefoothree",
+      "render-result-1": "onefoothree"
+    }
+  },
+  "tests/test_templating.py::TestTemplating::test_array_set_item": {
+    "recorded-date": "16-05-2025, 19:13:41",
+    "recorded-content": {
+      "render-result-1-cli": "twoonefoothree",
+      "render-result-1": "twoonefoothree"
+    }
+  },
+  "tests/test_templating.py::TestTemplating::test_dict_assign_quoted_item_via_brackets": {
+    "recorded-date": "28-05-2025, 11:06:39",
+    "recorded-content": {
+      "render-result-1-cli": "{bar=foo}",
+      "render-result-1": "{bar=foo}"
+    }
+  },
+  "tests/test_templating.py::TestTemplating::test_dict_assign_nested_items_via_brackets": {
+    "recorded-date": "29-05-2025, 11:56:01",
+    "recorded-content": {
+      "render-result-1-cli": "foo",
+      "render-result-1": "foo"
+    }
+  },
+  "tests/test_templating.py::TestTemplating::test_dict_assign_var_not_found": {
+    "recorded-date": "29-05-2025, 12:07:42",
+    "recorded-content": {
+      "render-result-1-cli": "{$key=foo}",
+      "render-result-1": "{=foo}"
+    }
   }
 }