improve Assignment operation

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.9: Small improvement of `Assignment`, introducing `AssignmentTerm` to help future expandability.
 * 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

airspeed/operators.py

@@ -1075,69 +1075,78 @@ def evaluate_raw(self, stream, namespace, loader):
             self.else_block.evaluate(stream, namespace, loader)
 
 
+class DotDictAccessTerm(_Element):
+    START = re.compile(r"\.([\w-]+)(.*)$", re.S + re.I)
+    value: str
+
+    def parse(self):
+        (self.value,) = self.identity_match(self.START)
+
+    def calculate(self, namespace, loader):
+        return self.value
+
+
+class BracketedValueTerm(_Element):
+    START = re.compile(r"(\[)\s*(.+)$", re.S + re.I)
+    END = re.compile(r"\s*](.*)$", re.S)
+    value: Value
+
+    def parse(self):
+        self.identity_match(self.START)
+        self.value = self.require_next_element(Value, "value")
+        self.require_match(self.END, "]")
+
+    def calculate(self, namespace, loader):
+        return self.value.calculate(namespace, loader)
+
+
+class AssignmentTerm(_Element):
+    term: DotDictAccessTerm | BracketedValueTerm
+
+    def parse(self):
+        self.term = self.next_element([DotDictAccessTerm, BracketedValueTerm])
+
+    def calculate(self, namespace, loader):
+        return self.term.calculate(namespace, loader)
+
+
 # This can't deal with assignments like
 # set($one.two().three = something)
 # yet
 class Assignment(_Element):
-    terms: list
-    START = re.compile(
-        # 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,
-    )
+    START = re.compile(r"\s*\(\s*\$(\w+)(.*)$", 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)
 
+    root_term: str
+    terms: list[AssignmentTerm]
+    value: Expression
+
     def parse(self):
-        self.identity_match(self.START)
+        (self.root_term,) = 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
+        while True:
+            try:
+                self.terms.append(self.next_element(AssignmentTerm))
+            except NoMatch:
                 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, ")")
 
     def evaluate_raw(self, stream, namespace, loader):
         val = self.value.calculate(namespace, loader)
-        if len(self.terms) == 1:
-            namespace.set_inherited(self.terms[0], val)
-        else:
-            cur = namespace
-            for term in self.terms[:-1]:
-                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
+
+        if not self.terms:
+            return namespace.set_inherited(self.root_term, val)
+
+        cur = namespace[self.root_term]
+        for term in self.terms[:-1]:
+            cur = cur[term.calculate(namespace, loader)]
+
+        cur[self.terms[-1].calculate(namespace, loader)] = val
 
 
 class EvaluateDirective(_Element):

setup.py

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