From 792feaacaec9e87ac9c4a86f85b5a2f665b99612 Mon Sep 17 00:00:00 2001
From: Cesar Alvernaz <cesar.alvernaz@gmail.com>
Date: Sat, 30 Sep 2023 09:59:28 +0100
Subject: [PATCH] array notation dict access with string literals, include
 upstream patches (#14)

---
 README.md                           |  1 +
 airspeed/operators.py               | 10 +++++-----
 setup.py                            |  2 +-
 tests/test_templating.py            |  3 +++
 tests/test_templating.snapshot.json |  7 +++++++
 5 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 108a8f4..c3c8dbb 100644
--- a/README.md
+++ b/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.3: array notation for dicts using string literals and merge upstream patches
 * v0.6.2: add support to contains and toString functions
 * v0.6.1: improve handling of multi-line dict expressions
 * v0.6.0: add initial setup for snapshot testing against AWS and Java VTL; enhance AWS parity
diff --git a/airspeed/operators.py b/airspeed/operators.py
index 2e14d0f..53e53d3 100644
--- a/airspeed/operators.py
+++ b/airspeed/operators.py
@@ -11,6 +11,7 @@
 
 LOG = logging.getLogger(__name__)
 
+
 # A dict that maps classes to dicts of additional methods.
 # This allows support for methods that are available in Java-based Velocity
 # implementations, e.g., .size() of a list or .length() of a string.
@@ -590,7 +591,7 @@ def calculate(self, namespace, loader):
 
 
 class NameOrCall(_Element):
-    NAME = re.compile(r"([a-zA-Z0-9_]+)(.*)$", re.S)
+    NAME = re.compile(r"([a-zA-Z0-9_-]+)(.*)$", re.S)
     parameters = None
     index = None
 
@@ -689,10 +690,8 @@ class VariableExpression(_Element):
 
     def parse(self):
         self.part = self.next_element(NameOrCall)
-        try:
+        with contextlib.suppress(NoMatch):
             self.subexpression = self.next_element(SubExpression)
-        except NoMatch:
-            pass
 
     def calculate(self, namespace, loader, global_namespace=None):
         if global_namespace is None:
@@ -732,6 +731,7 @@ def parse(self):
             (
                 FormalReference,
                 IntegerLiteral,
+                StringLiteral,
                 InterpolatedStringLiteral,
                 ParenthesizedExpression,
             ),
@@ -1047,7 +1047,7 @@ def evaluate_raw(self, stream, namespace, loader):
 # yet
 class Assignment(_Element):
     START = re.compile(
-        r"\s*\(\s*\$([a-z_][a-z0-9_]*(?:\.[a-z_][a-z0-9_]*)*)\s*=\s*(.*)$", re.S + re.I
+        r"\s*\(\s*\$(\w*(?:\.[\w-]+|\[\"\$\w+\"\]*)*)\s*=\s*(.*)$", re.S + re.I
     )
     END = re.compile(r"\s*\)(?:[ \t]*\r?\n)?(.*)$", re.S + re.M)
 
diff --git a/setup.py b/setup.py
index 5a805b4..a8a5149 100755
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setup(
     name="airspeed-ext",
-    version="0.6.2",
+    version="0.6.3",
     description=(
         "Airspeed is a powerful and easy-to-use templating engine "
         "for Python that aims for a high level of compatibility "
diff --git a/tests/test_templating.py b/tests/test_templating.py
index 780b43e..3641b20 100644
--- a/tests/test_templating.py
+++ b/tests/test_templating.py
@@ -729,6 +729,9 @@ def test_array_notation_dot(self, test_render):
     def test_array_notation_dict_index(self, test_render):
         test_render('$a["foo"]', {"a": {"foo": "bar"}})
 
+    def test_array_notation_dict_index_with_string_literal(self, test_render):
+        test_render("$a.b.c['foo:bar']", {"a": {"b": {"c": {"foo:bar": "baz"}}}})
+
     def test_array_notation_variable_index(self, test_render):
         test_render("#set($i = 1)$a[ $i ]", {"a": ["foo", "bar"]}, skip_cli=True)
 
diff --git a/tests/test_templating.snapshot.json b/tests/test_templating.snapshot.json
index 6f4474a..5c19728 100644
--- a/tests/test_templating.snapshot.json
+++ b/tests/test_templating.snapshot.json
@@ -1033,5 +1033,12 @@
       "render-result-1-cli": " yes!",
       "render-result-1": " yes!"
     }
+  },
+  "tests/test_templating.py::TestTemplating::test_array_notation_dict_index_with_string_literal": {
+    "recorded-date": "28-09-2023, 20:46:17",
+    "recorded-content": {
+      "render-result-1-cli": "$a.b.c['foo:bar']",
+      "render-result-1": "baz"
+    }
   }
 }