Repository: smithy-lang/smithy
Pull request #2033 (2023-11-07 13:59:23)

Add convenience method to get IntEnumShape from GenerateIntEnumDirective
*Description of changes:*
Adds convenience method to get an IntEnumShape from the `GenerateIntEnumDirective`.

This provides parity with the `GenerateEnumDirective` that has an `expectEnumShape` method [see here](https://github.com/smithy-lang/smithy/blob/3ea52156715f7e30de026763021ef78447bf965d/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateEnumDirective.java#L79)

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.



Patch:
diff --git a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateIntEnumDirective.java b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateIntEnumDirective.java
index 1ab87b7f7..dc3518af6 100644
--- a/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateIntEnumDirective.java
+++ b/smithy-codegen-core/src/main/java/software/amazon/smithy/codegen/core/directed/GenerateIntEnumDirective.java
@@ -16,6 +16,8 @@
 package software.amazon.smithy.codegen.core.directed;
 
 import software.amazon.smithy.codegen.core.CodegenContext;
+import software.amazon.smithy.model.node.ExpectationNotMetException;
+import software.amazon.smithy.model.shapes.IntEnumShape;
 import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.shapes.Shape;
 
@@ -29,6 +31,19 @@ import software.amazon.smithy.model.shapes.Shape;
 public final class GenerateIntEnumDirective<C extends CodegenContext<S, ?, ?>, S> extends ShapeDirective<Shape, C, S> {
 
     GenerateIntEnumDirective(C context, ServiceShape service, Shape shape) {
-        super(context, service, shape);
+        super(context, service, validateShape(shape));
+    }
+
+    private static Shape validateShape(Shape shape) {
+        if (shape.isIntEnumShape()) {
+            return shape;
+        }
+        throw new IllegalArgumentException("GenerateIntEnum requires an IntEnum shape");
+    }
+
+
+    public IntEnumShape expectIntEnumShape() {
+        return shape().asIntEnumShape().orElseThrow(() -> new ExpectationNotMetException(
+                "Expected an IntEnum shape, but found " + shape(), shape()));
     }
 }


Test patch:
diff --git a/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/GenerateIntEnumDirectiveTest.java b/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/GenerateIntEnumDirectiveTest.java
new file mode 100644
index 000000000..2d582fbab
--- /dev/null
+++ b/smithy-codegen-core/src/test/java/software/amazon/smithy/codegen/core/directed/GenerateIntEnumDirectiveTest.java
@@ -0,0 +1,16 @@
+package software.amazon.smithy.codegen.core.directed;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import software.amazon.smithy.model.shapes.BooleanShape;
+
+public class GenerateIntEnumDirectiveTest {
+    @Test
+    public void validatesShapeType() {
+        BooleanShape shape = BooleanShape.builder().id("smithy.example#Foo").build();
+
+        Assertions.assertThrows(IllegalArgumentException.class, () -> {
+            new GenerateIntEnumDirective<TestContext, TestSettings>(null, null, shape);
+        });
+    }
+}