Jmespath custom functions

.changes/next-release/feature-dedce719f4dc02f94249c5b65ce4d54e508f1232.json

@@ -0,0 +1,7 @@
+{
+  "type": "feature",
+  "description": "Added a generic evaluator/interpreter for JMESPath expressions.",
+  "pull_requests": [
+    "[#2878](https://github.com/smithy-lang/smithy/pull/2878)"
+  ]
+}

smithy-jmespath-tests/src/main/java/software/amazon/smithy/jmespath/tests/ComplianceTestRunner.java

@@ -11,16 +11,20 @@
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.function.BiPredicate;
 import java.util.stream.Stream;
 import software.amazon.smithy.jmespath.JmespathException;
 import software.amazon.smithy.jmespath.JmespathExceptionType;
 import software.amazon.smithy.jmespath.JmespathExpression;
 import software.amazon.smithy.jmespath.RuntimeType;
+import software.amazon.smithy.jmespath.evaluation.EvaluationUtils;
 import software.amazon.smithy.jmespath.evaluation.Evaluator;
+import software.amazon.smithy.jmespath.evaluation.JmespathAbstractRuntime;
 import software.amazon.smithy.jmespath.evaluation.JmespathRuntime;
+import software.amazon.smithy.jmespath.type.Type;
 import software.amazon.smithy.utils.IoUtils;
 
-public class ComplianceTestRunner<T> {
+public class ComplianceTestRunner<T, A extends Type> {
     private static final String DEFAULT_TEST_CASE_LOCATION = "compliance";
     private static final String SUBJECT_MEMBER = "given";
     private static final String CASES_MEMBER = "cases";
@@ -30,19 +34,25 @@ public class ComplianceTestRunner<T> {
     private static final String ERROR_MEMBER = "error";
     private static final String BENCH_MEMBER = "bench";
     private final JmespathRuntime<T> runtime;
-    private final List<TestCase<T>> testCases = new ArrayList<>();
+    private final JmespathAbstractRuntime<A> abstractRuntime;
+    private final List<TestCase<T, A>> testCases = new ArrayList<>();
 
-    private ComplianceTestRunner(JmespathRuntime<T> runtime) {
+    private ComplianceTestRunner(JmespathRuntime<T> runtime, JmespathAbstractRuntime<A> abstractRuntime) {
         this.runtime = runtime;
+        this.abstractRuntime = abstractRuntime;
     }
 
     public static <T> Stream<Object[]> defaultParameterizedTestSource(JmespathRuntime<T> runtime) {
-        ComplianceTestRunner<T> runner = new ComplianceTestRunner<>(runtime);
+        return defaultParameterizedTestSource(runtime, null);
+    }
+
+    public static <T, A extends Type> Stream<Object[]> defaultParameterizedTestSource(JmespathRuntime<T> runtime, JmespathAbstractRuntime<A> abstractRuntime) {
+        ComplianceTestRunner<T, A> runner = new ComplianceTestRunner<>(runtime, abstractRuntime);
         URL manifest = ComplianceTestRunner.class.getResource(DEFAULT_TEST_CASE_LOCATION + "/MANIFEST");
         try (var reader = new BufferedReader(new InputStreamReader(manifest.openStream(), StandardCharsets.UTF_8))) {
             reader.lines().forEach(line -> {
                 var url = ComplianceTestRunner.class.getResource(DEFAULT_TEST_CASE_LOCATION + "/" + line.trim());
-                runner.testCases.addAll(TestCase.from(url, runtime));
+                runner.testCases.addAll(TestCase.from(url, runtime, abstractRuntime));
             });
         } catch (IOException e) {
             throw new RuntimeException(e);
@@ -54,8 +64,13 @@ public Stream<Object[]> parameterizedTestSource() {
         return testCases.stream().map(testCase -> new Object[] {testCase.name(), testCase});
     }
 
-    private record TestCase<T>(
+    public <A> Stream<Object[]> parameterizedAbstractTestSource(JmespathRuntime<A> abstractRuntime, BiPredicate<T, A> abstractPredicate) {
+        return testCases.stream().map(testCase -> new Object[] {testCase.name(), (Runnable)() -> testCase.abstractRun(abstractRuntime, abstractPredicate)});
+    }
+
+    private record TestCase<T, A extends Type>(
             JmespathRuntime<T> runtime,
+            JmespathAbstractRuntime<A> abstractRuntime,
             String testSuite,
             String comment,
             T given,
@@ -64,10 +79,10 @@ private record TestCase<T>(
             JmespathExceptionType expectedError,
             String benchmark)
             implements Runnable {
-        public static <T> List<TestCase<T>> from(URL url, JmespathRuntime<T> runtime) {
+        public static <T, A extends Type> List<TestCase<T, A>> from(URL url, JmespathRuntime<T> runtime, JmespathAbstractRuntime<A> abstractRuntime) {
             var path = url.getPath();
             var testSuiteName = path.substring(path.lastIndexOf('/') + 1, path.lastIndexOf('.'));
-            var testCases = new ArrayList<TestCase<T>>();
+            var testCases = new ArrayList<TestCase<T, A>>();
             String text = IoUtils.readUtf8Url(url);
             T tests = JmespathExpression.parseJson(text, runtime);
 
@@ -92,6 +107,7 @@ public static <T> List<TestCase<T>> from(URL url, JmespathRuntime<T> runtime) {
 
                     var benchmark = valueAsString(runtime, testCase, BENCH_MEMBER);
                     testCases.add(new TestCase<>(runtime,
+                            abstractRuntime,
                             testSuiteName,
                             comment,
                             given,
@@ -122,9 +138,10 @@ private String name() {
         public void run() {
             try {
                 var parsed = JmespathExpression.parse(expression);
-                var result = new Evaluator<>(given, runtime).visit(parsed);
+                var result = parsed.evaluate(given, runtime);
                 if (benchmark != null) {
                     // Benchmarks don't include expected results or errors
+                    // TODO: Could still run these?
                     return;
                 }
                 if (expectedError != null) {
@@ -138,7 +155,34 @@ public void run() {
                                 + "Actual:    " + runtime.toString(result) + "\n"
                                 + "For query: " + expression + "\n");
                     }
+
+                    if (abstractRuntime != null) {
+                        var abstractedGiven = EvaluationUtils.convert(runtime, given, abstractRuntime);
+                        var abstractResult = parsed.evaluate(abstractedGiven, abstractRuntime);
+
+                        if (!abstractResult.isInstance(result, runtime)) {
+                            parsed.evaluate(abstractedGiven, abstractRuntime);
+                            abstractResult.isInstance(result, runtime);
+                            throw new AssertionError("Expected " + result + " to be an instance of " + abstractResult + ".\n"
+                                    + "For query: " + expression + "\n");
+                        }
+                    }
+                }
+            } catch (JmespathException e) {
+                if (!e.getType().equals(expectedError)) {
+                    throw new AssertionError("Expected error does not match actual error. \n"
+                            + "Expected: " + (expectedError != null ? expectedError : "(no error)") + "\n"
+                            + "Actual: " + e.getType() + " - " + e.getMessage() + "\n"
+                            + "For query: " + expression + "\n", e);
                 }
+            }
+        }
+
+        public <A> void abstractRun(JmespathRuntime<A> abstractRuntime, BiPredicate<T, A> abstractPredicate) {
+            try {
+                var parsed = JmespathExpression.parse(expression);
+                var result = new Evaluator<>(given, runtime).visit(parsed);
+
             } catch (JmespathException e) {
                 if (!e.getType().equals(expectedError)) {
                     throw new AssertionError("Expected error does not match actual error. \n"

smithy-jmespath-tests/src/test/java/software/amazon/smithy/jmespath/tests/LiteralExpressionJmespathRuntimeComplianceTests.java

@@ -8,6 +8,7 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
 import software.amazon.smithy.jmespath.LiteralExpressionJmespathRuntime;
+import software.amazon.smithy.jmespath.type.TypeJmespathRuntime;
 
 public class LiteralExpressionJmespathRuntimeComplianceTests {
     @ParameterizedTest(name = "{0}")
@@ -17,6 +18,8 @@ public void testRunner(String filename, Runnable callable) throws Exception {
     }
 
     public static Stream<?> source() {
-        return ComplianceTestRunner.defaultParameterizedTestSource(LiteralExpressionJmespathRuntime.INSTANCE);
+        return ComplianceTestRunner.defaultParameterizedTestSource(
+                LiteralExpressionJmespathRuntime.INSTANCE,
+                new TypeJmespathRuntime());
     }
 }

smithy-jmespath/src/main/java/software/amazon/smithy/jmespath/JmespathExpression.java

@@ -7,7 +7,9 @@
 import java.util.Set;
 import java.util.TreeSet;
 import software.amazon.smithy.jmespath.ast.LiteralExpression;
+import software.amazon.smithy.jmespath.evaluation.AbstractEvaluator;
 import software.amazon.smithy.jmespath.evaluation.Evaluator;
+import software.amazon.smithy.jmespath.evaluation.JmespathAbstractRuntime;
 import software.amazon.smithy.jmespath.evaluation.JmespathRuntime;
 
 /**
@@ -42,7 +44,7 @@ public static JmespathExpression parse(String text) {
      * @return Returns the parsed expression.
      * @throws JmespathException if the expression is invalid.
      */
-    public static <T> JmespathExpression parse(String text, JmespathRuntime<T> runtime) {
+    public static <T> JmespathExpression parse(String text, JmespathAbstractRuntime<T> runtime) {
         return Parser.parse(text, runtime);
     }
 
@@ -54,7 +56,7 @@ public static <T> JmespathExpression parse(String text, JmespathRuntime<T> runti
      * @return Returns the parsed JSON value.
      * @throws JmespathException if the text is invalid.
      */
-    public static <T> T parseJson(String text, JmespathRuntime<T> runtime) {
+    public static <T> T parseJson(String text, JmespathAbstractRuntime<T> runtime) {
         Lexer<T> lexer = new Lexer<T>(text, runtime);
         return lexer.parseJsonValue();
     }
@@ -129,4 +131,12 @@ public LiteralExpression evaluate(LiteralExpression currentNode) {
     public <T> T evaluate(T currentNode, JmespathRuntime<T> runtime) {
         return new Evaluator<>(currentNode, runtime).visit(this);
     }
+
+    public <T> T evaluate(T currentNode, JmespathAbstractRuntime<T> runtime) {
+        if (runtime instanceof JmespathRuntime) {
+            return evaluate(currentNode, (JmespathRuntime<T>)runtime);
+        } else {
+            return new AbstractEvaluator<>(currentNode, runtime).visit(this);
+        }
+    }
 }

smithy-jmespath/src/main/java/software/amazon/smithy/jmespath/JmespathExtension.java

@@ -0,0 +1,11 @@
+package software.amazon.smithy.jmespath;
+
+
+import software.amazon.smithy.jmespath.evaluation.Function;
+
+import java.util.List;
+
+public interface JmespathExtension {
+
+    <T> List<Function<T>> getFunctions();
+}

smithy-jmespath/src/main/java/software/amazon/smithy/jmespath/JmespathQuery.java

@@ -0,0 +1,11 @@
+package software.amazon.smithy.jmespath;
+
+
+import software.amazon.smithy.jmespath.evaluation.JmespathRuntime;
+
+import java.util.function.Function;
+
+public interface JmespathQuery<T> extends Function<T, T> {
+
+    JmespathRuntime<T> runtime();
+}

smithy-jmespath/src/main/java/software/amazon/smithy/jmespath/Lexer.java

@@ -9,13 +9,14 @@
 import java.util.Objects;
 import java.util.function.Predicate;
 import software.amazon.smithy.jmespath.ast.LiteralExpression;
+import software.amazon.smithy.jmespath.evaluation.JmespathAbstractRuntime;
 import software.amazon.smithy.jmespath.evaluation.JmespathRuntime;
 
 final class Lexer<T> {
 
     private static final int MAX_NESTING_LEVEL = 50;
 
-    private final JmespathRuntime<T> runtime;
+    private final JmespathAbstractRuntime<T> runtime;
     private final String expression;
     private final int length;
     private int position = 0;
@@ -25,7 +26,7 @@ final class Lexer<T> {
     private final List<Token> tokens = new ArrayList<>();
     private boolean currentlyParsingLiteral;
 
-    Lexer(String expression, JmespathRuntime<T> runtime) {
+    Lexer(String expression, JmespathAbstractRuntime<T> runtime) {
         this.runtime = Objects.requireNonNull(runtime, "runtime must not be null");
         this.expression = Objects.requireNonNull(expression, "expression must not be null");
         this.length = expression.length();
@@ -35,7 +36,7 @@ static TokenIterator tokenize(String expression) {
         return tokenize(expression, LiteralExpressionJmespathRuntime.INSTANCE);
     }
 
-    static <T> TokenIterator tokenize(String expression, JmespathRuntime<T> runtime) {
+    static <T> TokenIterator tokenize(String expression, JmespathAbstractRuntime<T> runtime) {
         return new Lexer<>(expression, runtime).doTokenize();
     }
 

smithy-jmespath/src/main/java/software/amazon/smithy/jmespath/LiteralExpressionJmespathRuntime.java

@@ -88,6 +88,11 @@ public LiteralExpression element(LiteralExpression array, int index) {
         return LiteralExpression.from(array.expectArrayValue().get(index));
     }
 
+    @Override
+    public LiteralExpression abstractElement(LiteralExpression array, LiteralExpression index) {
+        return element(array, index.expectNumberValue().intValue());
+    }
+
     @Override
     public Iterable<LiteralExpression> asIterable(LiteralExpression array) {
         switch (array.getType()) {
@@ -109,17 +114,19 @@ private static final class ArrayLiteralExpressionBuilder implements ArrayBuilder
         private final List<Object> result = new ArrayList<>();
 
         @Override
-        public void add(LiteralExpression value) {
+        public ArrayLiteralExpressionBuilder add(LiteralExpression value) {
             result.add(value.getValue());
+            return this;
         }
 
         @Override
-        public void addAll(LiteralExpression array) {
+        public ArrayLiteralExpressionBuilder addAll(LiteralExpression array) {
             if (array.isArrayValue()) {
                 result.addAll(array.expectArrayValue());
             } else {
                 result.addAll(array.expectObjectValue().keySet());
             }
+            return this;
         }
 
         @Override
@@ -146,13 +153,15 @@ private static final class ObjectLiteralExpressionBuilder implements ObjectBuild
         private final Map<String, Object> result = new HashMap<>();
 
         @Override
-        public void put(LiteralExpression key, LiteralExpression value) {
+        public ObjectLiteralExpressionBuilder put(LiteralExpression key, LiteralExpression value) {
             result.put(key.expectStringValue(), value.getValue());
+            return this;
         }
 
         @Override
-        public void putAll(LiteralExpression object) {
+        public ObjectLiteralExpressionBuilder putAll(LiteralExpression object) {
             result.putAll(object.expectObjectValue());
+            return this;
         }
 
         @Override

smithy-jmespath/src/main/java/software/amazon/smithy/jmespath/Parser.java

@@ -27,6 +27,7 @@
 import software.amazon.smithy.jmespath.ast.ProjectionExpression;
 import software.amazon.smithy.jmespath.ast.SliceExpression;
 import software.amazon.smithy.jmespath.ast.Subexpression;
+import software.amazon.smithy.jmespath.evaluation.JmespathAbstractRuntime;
 import software.amazon.smithy.jmespath.evaluation.JmespathRuntime;
 
 /**
@@ -75,12 +76,12 @@ final class Parser {
     private final String expression;
     private final TokenIterator iterator;
 
-    private Parser(String expression, JmespathRuntime<?> runtime) {
+    private Parser(String expression, JmespathAbstractRuntime<?> runtime) {
         this.expression = expression;
         iterator = Lexer.tokenize(expression, runtime);
     }
 
-    static JmespathExpression parse(String expression, JmespathRuntime<?> runtime) {
+    static JmespathExpression parse(String expression, JmespathAbstractRuntime<?> runtime) {
         Parser parser = new Parser(expression, runtime);
         JmespathExpression result = parser.expression(0);
         parser.iterator.expect(TokenType.EOF);

smithy-jmespath/src/main/java/software/amazon/smithy/jmespath/RuntimeType.java

@@ -4,6 +4,7 @@
  */
 package software.amazon.smithy.jmespath;
 
+import java.util.EnumSet;
 import java.util.Locale;
 import software.amazon.smithy.jmespath.ast.ComparatorType;
 import software.amazon.smithy.jmespath.ast.LiteralExpression;
@@ -150,4 +151,8 @@ public abstract LiteralExpression compare(
             LiteralExpression right,
             ComparatorType comparator
     );
+
+    public static EnumSet<RuntimeType> valueTypes() {
+        return EnumSet.complementOf(EnumSet.of(EXPRESSION, ANY));
+    }
 }