Use Smithy's JMESPath evaluator

.github/workflows/ci.yml

@@ -19,6 +19,16 @@ jobs:
         os: [macos-latest, ubuntu-latest]
 
     steps:
+      # HACK TEMPORARY WORKAROUND UNTIL https://github.com/smithy-lang/smithy/pull/2878 IS RELEASED
+      - uses: actions/checkout@v6
+        with:
+          repository: 'robin-aws/smithy'
+          path: 'smithy'
+          ref: 'smithy-jmespath-evaluator'
+      - working-directory: smithy
+        run: ./gradlew publishToMavenLocal
+      # END HACK
+
       - uses: actions/checkout@v6
       - name: Set up JDK ${{ matrix.java }}
         uses: actions/setup-java@v5

codegen/codegen-core/src/main/java/software/amazon/smithy/java/codegen/generators/EnumGenerator.java

@@ -20,6 +20,8 @@
 import software.amazon.smithy.java.codegen.sections.EnumVariantSection;
 import software.amazon.smithy.java.codegen.writer.JavaWriter;
 import software.amazon.smithy.java.core.schema.SerializableShape;
+import software.amazon.smithy.java.core.schema.SmithyEnum;
+import software.amazon.smithy.java.core.schema.SmithyIntEnum;
 import software.amazon.smithy.java.core.serde.ShapeDeserializer;
 import software.amazon.smithy.java.core.serde.ShapeSerializer;
 import software.amazon.smithy.model.Model;
@@ -44,7 +46,7 @@ public void accept(T directive) {
         directive.context().writerDelegator().useShapeWriter(shape, writer -> {
             writer.pushState(new ClassSection(shape));
             var template = """
-                    public sealed interface ${shape:T} extends ${serializableShape:T} {
+                    public sealed interface ${shape:T} extends ${enum:T}, ${serializableShape:T} {
                         ${staticImpls:C|}
 
                         ${schema:C|}
@@ -68,6 +70,7 @@ default void serialize(${shapeSerializer:T} serializer) {
             var shapeSymbol = directive.symbolProvider().toSymbol(shape);
             writer.putContext("shape", shapeSymbol);
             writer.putContext("serializableShape", SerializableShape.class);
+            writer.putContext("enum", shape.isEnumShape() ? SmithyEnum.class : SmithyIntEnum.class);
             writer.putContext("shapeSerializer", ShapeSerializer.class);
             var valueSymbol = shapeSymbol.expectProperty(SymbolProperties.ENUM_VALUE_TYPE);
             writer.putContext("value", valueSymbol);

codegen/plugins/types-codegen/src/test/resources/software/amazon/smithy/java/codegen/types/test-cases/enums/expected/EnumType.java

@@ -7,14 +7,15 @@
 import software.amazon.smithy.java.core.schema.Schema;
 import software.amazon.smithy.java.core.schema.SerializableShape;
 import software.amazon.smithy.java.core.schema.ShapeBuilder;
+import software.amazon.smithy.java.core.schema.SmithyEnum;
 import software.amazon.smithy.java.core.serde.ShapeDeserializer;
 import software.amazon.smithy.java.core.serde.ShapeSerializer;
 import software.amazon.smithy.java.core.serde.ToStringSerializer;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.utils.SmithyGenerated;
 
 @SmithyGenerated
-public sealed interface EnumType extends SerializableShape {
+public sealed interface EnumType extends SmithyEnum, SerializableShape {
     EnumType OPTION_ONE = new OptionOneType();
     EnumType OPTION_TWO = new OptionTwoType();
     List<EnumType> $TYPES = List.of(OPTION_ONE, OPTION_TWO);

codegen/plugins/types-codegen/src/test/resources/software/amazon/smithy/java/codegen/types/test-cases/enums/expected/IntEnumType.java

@@ -6,14 +6,15 @@
 import software.amazon.smithy.java.core.schema.Schema;
 import software.amazon.smithy.java.core.schema.SerializableShape;
 import software.amazon.smithy.java.core.schema.ShapeBuilder;
+import software.amazon.smithy.java.core.schema.SmithyIntEnum;
 import software.amazon.smithy.java.core.serde.ShapeDeserializer;
 import software.amazon.smithy.java.core.serde.ShapeSerializer;
 import software.amazon.smithy.java.core.serde.ToStringSerializer;
 import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.utils.SmithyGenerated;
 
 @SmithyGenerated
-public sealed interface IntEnumType extends SerializableShape {
+public sealed interface IntEnumType extends SmithyIntEnum, SerializableShape {
     IntEnumType FIRST = new FirstType();
     IntEnumType SECOND = new SecondType();
     IntEnumType FIFTH = new FifthType();

core/src/main/java/software/amazon/smithy/java/core/schema/SmithyEnum.java

@@ -0,0 +1,11 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.java.core.schema;
+
+public interface SmithyEnum {
+
+    String getValue();
+}

core/src/main/java/software/amazon/smithy/java/core/schema/SmithyIntEnum.java

@@ -0,0 +1,11 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.java.core.schema;
+
+public interface SmithyIntEnum {
+
+    int getValue();
+}

gradle/libs.versions.toml

@@ -32,6 +32,7 @@ smithy-aws-protocol-tests = { module = "software.amazon.smithy:smithy-aws-protoc
 smithy-protocol-tests = { module = "software.amazon.smithy:smithy-protocol-tests", version.ref = "smithy" }
 smithy-validation-model = { module = "software.amazon.smithy:smithy-validation-model", version.ref = "smithy" }
 smithy-jmespath = { module = "software.amazon.smithy:smithy-jmespath", version.ref = "smithy" }
+smithy-jmespath-tests = { module = "software.amazon.smithy:smithy-jmespath-tests", version.ref = "smithy" }
 smithy-waiters = { module = "software.amazon.smithy:smithy-waiters", version.ref = "smithy" }
 smithy-utils = { module = "software.amazon.smithy:smithy-utils", version.ref = "smithy" }
 smithy-traitcodegen = { module = "software.amazon.smithy:smithy-trait-codegen", version.ref = "smithy" }

jmespath/build.gradle.kts

@@ -10,4 +10,5 @@ extra["moduleName"] = "software.amazon.smithy.java.jmespath"
 dependencies {
     api(project(":core"))
     api(libs.smithy.jmespath)
+    testImplementation(libs.smithy.jmespath.tests)
 }

jmespath/src/main/java/software/amazon/smithy/java/jmespath/DocumentJmespathRuntime.java

@@ -0,0 +1,165 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.java.jmespath;
+
+import java.math.BigDecimal;
+import java.math.BigInteger;
+import software.amazon.smithy.java.core.serde.SerializationException;
+import software.amazon.smithy.java.core.serde.document.Document;
+import software.amazon.smithy.jmespath.JmespathException;
+import software.amazon.smithy.jmespath.JmespathExceptionType;
+import software.amazon.smithy.jmespath.RuntimeType;
+import software.amazon.smithy.jmespath.evaluation.EvaluationUtils;
+import software.amazon.smithy.jmespath.evaluation.JmespathRuntime;
+import software.amazon.smithy.jmespath.evaluation.ListArrayBuilder;
+import software.amazon.smithy.jmespath.evaluation.MapObjectBuilder;
+import software.amazon.smithy.jmespath.evaluation.MappingIterable;
+import software.amazon.smithy.jmespath.evaluation.NumberType;
+import software.amazon.smithy.model.shapes.ShapeType;
+
+// TODO: default equal isn't correct, need to normalize number types
+public class DocumentJmespathRuntime implements JmespathRuntime<Document> {
+
+    public static final DocumentJmespathRuntime INSTANCE = new DocumentJmespathRuntime();
+
+    @Override
+    public RuntimeType typeOf(Document document) {
+        if (document == null) {
+            return RuntimeType.NULL;
+        }
+        return switch (document.type()) {
+            case BOOLEAN ->
+                RuntimeType.BOOLEAN;
+            case BLOB, ENUM, STRING ->
+                RuntimeType.STRING;
+            case BYTE, SHORT, INTEGER, INT_ENUM, LONG, FLOAT, DOUBLE, BIG_DECIMAL, BIG_INTEGER, TIMESTAMP ->
+                RuntimeType.NUMBER;
+            case LIST, SET ->
+                RuntimeType.ARRAY;
+            case MAP, STRUCTURE, UNION ->
+                RuntimeType.OBJECT;
+            default -> throw new IllegalArgumentException("Unknown runtime type: " + document.type());
+        };
+    }
+
+    @Override
+    public Document createNull() {
+        return null;
+    }
+
+    @Override
+    public Document createBoolean(boolean b) {
+        return Document.of(b);
+    }
+
+    @Override
+    public boolean asBoolean(Document document) {
+        return document.asBoolean();
+    }
+
+    @Override
+    public Document createString(String s) {
+        return Document.of(s);
+    }
+
+    @Override
+    public String asString(Document document) {
+        if (document.isType(ShapeType.BLOB)) {
+            return JMESPathDocumentUtils.asString(document.asBlob());
+        } else {
+            try {
+                return document.asString();
+            } catch (SerializationException e) {
+                throw new JmespathException(JmespathExceptionType.INVALID_TYPE, "Not a string: " + document, e);
+            }
+        }
+    }
+
+    @Override
+    public Document createNumber(Number number) {
+        return switch (EvaluationUtils.numberType(number)) {
+            case BYTE -> Document.of(number.byteValue());
+            case SHORT -> Document.of(number.shortValue());
+            case INTEGER -> Document.of(number.intValue());
+            case LONG -> Document.of(number.longValue());
+            case FLOAT -> Document.of(number.floatValue());
+            case DOUBLE -> Document.of(number.doubleValue());
+            case BIG_INTEGER -> Document.of((BigInteger) number);
+            case BIG_DECIMAL -> Document.of((BigDecimal) number);
+        };
+    }
+
+    @Override
+    public NumberType numberType(Document document) {
+        return switch (document.type()) {
+            case BYTE -> NumberType.BYTE;
+            case SHORT -> NumberType.SHORT;
+            case INTEGER, INT_ENUM -> NumberType.INTEGER;
+            case LONG -> NumberType.LONG;
+            case FLOAT -> NumberType.FLOAT;
+            case DOUBLE -> NumberType.DOUBLE;
+            case BIG_DECIMAL, TIMESTAMP -> NumberType.BIG_DECIMAL;
+            case BIG_INTEGER -> NumberType.BIG_INTEGER;
+            default -> throw new IllegalArgumentException("Not a number: " + document);
+        };
+    }
+
+    @Override
+    public Number asNumber(Document document) {
+        if (document.isType(ShapeType.TIMESTAMP)) {
+            return JMESPathDocumentUtils.asBigDecimal(document.asTimestamp());
+        } else {
+            try {
+                return document.asNumber();
+            } catch (SerializationException e) {
+                throw new JmespathException(JmespathExceptionType.INVALID_TYPE, "Not a number: " + document, e);
+            }
+        }
+    }
+
+    @Override
+    public Number length(Document document) {
+        if (is(document, RuntimeType.STRING)) {
+            return EvaluationUtils.codePointCount(document.asString());
+        } else {
+            // This handles objects and arrays
+            return document.size();
+        }
+    }
+
+    @Override
+    public Document element(Document document, Document index) {
+        return document.asList().get(asNumber(index).intValue());
+    }
+
+    @Override
+    public Iterable<Document> asIterable(Document document) {
+        return switch (typeOf(document)) {
+            case ARRAY -> document.asList();
+            case OBJECT -> new MappingIterable<>(Document::of, document.asStringMap().keySet());
+            default -> throw new IllegalArgumentException("Not iterable: " + document);
+        };
+    }
+
+    @Override
+    public JmespathRuntime.ArrayBuilder<Document> arrayBuilder() {
+        return new ListArrayBuilder<>(this, Document::of);
+    }
+
+    @Override
+    public Document value(Document document, Document key) {
+        if (typeOf(document) == RuntimeType.OBJECT) {
+            return document.getMember(key.asString());
+        } else {
+            return createNull();
+        }
+    }
+
+    @Override
+    public JmespathRuntime.ObjectBuilder<Document> objectBuilder() {
+        return new MapObjectBuilder<>(this, Document::of);
+    }
+}