Expose event-type in stream interface

codegen/src/main/java/software/amazon/awssdk/codegen/poet/common/AbstractEnumClass.java

@@ -0,0 +1,115 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.codegen.poet.common;
+
+import static software.amazon.awssdk.codegen.poet.PoetUtils.createEnumBuilder;
+import static software.amazon.awssdk.codegen.poet.PoetUtils.toStringBuilder;
+
+import com.squareup.javapoet.ClassName;
+import com.squareup.javapoet.MethodSpec;
+import com.squareup.javapoet.ParameterizedTypeName;
+import com.squareup.javapoet.TypeSpec;
+import com.squareup.javapoet.TypeSpec.Builder;
+import java.util.Set;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import javax.lang.model.element.Modifier;
+import software.amazon.awssdk.codegen.model.intermediate.ShapeModel;
+import software.amazon.awssdk.codegen.poet.ClassSpec;
+
+public abstract class AbstractEnumClass implements ClassSpec {
+
+    private static final String VALUE = "value";
+    private static final String UNKNOWN_TO_SDK_VERSION = "UNKNOWN_TO_SDK_VERSION";
+    private final ShapeModel shape;
+
+    public AbstractEnumClass(ShapeModel shape) {
+        this.shape = shape;
+    }
+
+    @Override
+    public final TypeSpec poetSpec() {
+        Builder enumBuilder = createEnumBuilder(className())
+                .addField(String.class, VALUE, Modifier.PRIVATE, Modifier.FINAL)
+                .addMethod(toStringBuilder().addStatement("return $T.valueOf($N)", String.class, VALUE).build())
+                .addMethod(fromValueSpec())
+                .addMethod(knownValuesSpec())
+                .addMethod(createConstructor());
+
+        addDeprecated(enumBuilder);
+        addJavadoc(enumBuilder);
+        addEnumConstants(enumBuilder);
+
+        enumBuilder.addEnumConstant(UNKNOWN_TO_SDK_VERSION, TypeSpec.anonymousClassBuilder("null").build());
+
+        return enumBuilder.build();
+    }
+
+    protected final ShapeModel getShape() {
+        return shape;
+    }
+
+    protected abstract void addDeprecated(Builder enumBuilder);
+
+    protected abstract void addJavadoc(Builder enumBuilder);
+
+    protected abstract void addEnumConstants(Builder enumBuilder);
+
+    private MethodSpec createConstructor() {
+        return MethodSpec.constructorBuilder()
+                .addModifiers(Modifier.PRIVATE)
+                .addParameter(String.class, VALUE)
+                .addStatement("this.$1N = $1N", VALUE)
+                .build();
+    }
+
+    private MethodSpec fromValueSpec() {
+        return MethodSpec.methodBuilder("fromValue")
+                .returns(className())
+                .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
+                .addJavadoc("Use this in place of valueOf to convert the raw string returned by the service into the " +
+                        "enum value.\n\n" +
+                        "@param $N real value\n" +
+                        "@return $T corresponding to the value\n", VALUE, className())
+                .addParameter(String.class, VALUE)
+                .beginControlFlow("if ($N == null)", VALUE)
+                .addStatement("return null")
+                .endControlFlow()
+                .addStatement("return $1T.of($2T.values())\n" +
+                                ".filter(e -> e.toString().equals($3N))\n" +
+                                ".findFirst()\n" +
+                                ".orElse(UNKNOWN_TO_SDK_VERSION)",
+                        Stream.class,
+                        className(),
+                        VALUE)
+                .build();
+    }
+
+    private MethodSpec knownValuesSpec() {
+        return MethodSpec.methodBuilder("knownValues")
+                .returns(ParameterizedTypeName.get(ClassName.get(Set.class), className()))
+                .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
+                .addJavadoc("Use this in place of {@link #values()} to return a {@link Set} of all values known to the "
+                        + "SDK.\n"
+                        + "This will return all known enum values except {@link #$N}.\n\n"
+                        + "@return a {@link $T} of known {@link $T}s", UNKNOWN_TO_SDK_VERSION, Set.class, className())
+                .addStatement("return $T.of(values()).filter(v -> v != $N).collect($T.toSet())",
+                        Stream.class,
+                        UNKNOWN_TO_SDK_VERSION,
+                        Collectors.class)
+                .build();
+    }
+}

codegen/src/main/java/software/amazon/awssdk/codegen/poet/common/EnumClass.java

@@ -15,108 +15,39 @@
 
 package software.amazon.awssdk.codegen.poet.common;
 
-import static software.amazon.awssdk.codegen.poet.PoetUtils.addDeprecated;
-import static software.amazon.awssdk.codegen.poet.PoetUtils.addJavadoc;
-import static software.amazon.awssdk.codegen.poet.PoetUtils.createEnumBuilder;
-import static software.amazon.awssdk.codegen.poet.PoetUtils.toStringBuilder;
-
 import com.squareup.javapoet.ClassName;
-import com.squareup.javapoet.MethodSpec;
-import com.squareup.javapoet.ParameterizedTypeName;
 import com.squareup.javapoet.TypeSpec;
 import com.squareup.javapoet.TypeSpec.Builder;
-import java.util.Collections;
-import java.util.Set;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
-import javax.lang.model.element.Modifier;
 import software.amazon.awssdk.codegen.model.intermediate.ShapeModel;
-import software.amazon.awssdk.codegen.poet.ClassSpec;
-import software.amazon.awssdk.codegen.poet.StaticImport;
-
-public final class EnumClass implements ClassSpec {
+import software.amazon.awssdk.codegen.poet.PoetUtils;
 
-    private static final String VALUE = "value";
-    private static final String UNKNOWN_TO_SDK_VERSION = "UNKNOWN_TO_SDK_VERSION";
-    private final ShapeModel shape;
-    private final ClassName className;
+public final class EnumClass extends AbstractEnumClass {
+    private final String enumPackageName;
 
     public EnumClass(String enumPackage, ShapeModel shape) {
-        this.shape = shape;
-        this.className = ClassName.get(enumPackage, shape.getShapeName());
+        super(shape);
+        this.enumPackageName = enumPackage;
     }
 
     @Override
-    public TypeSpec poetSpec() {
-        Builder enumBuilder = createEnumBuilder(className)
-            .addField(String.class, VALUE, Modifier.PRIVATE, Modifier.FINAL)
-            .addMethod(toStringBuilder().addStatement("return $T.valueOf($N)", String.class, VALUE).build())
-            .addMethod(fromValueSpec())
-            .addMethod(knownValuesSpec())
-            .addMethod(createConstructor());
-
-        addDeprecated(enumBuilder::addAnnotation, shape);
-        addJavadoc(enumBuilder::addJavadoc, shape);
-
-        shape.getEnums().forEach(
-            e -> enumBuilder.addEnumConstant(e.getName(), TypeSpec.anonymousClassBuilder("$S", e.getValue()).build())
-        );
-        enumBuilder.addEnumConstant(UNKNOWN_TO_SDK_VERSION, TypeSpec.anonymousClassBuilder("null").build());
-
-        return enumBuilder.build();
+    protected void addDeprecated(Builder enumBuilder) {
+        PoetUtils.addDeprecated(enumBuilder::addAnnotation, getShape());
     }
 
     @Override
-    public ClassName className() {
-        return className;
+    protected void addJavadoc(Builder enumBuilder) {
+        PoetUtils.addJavadoc(enumBuilder::addJavadoc, getShape());
     }
 
     @Override
-    public Iterable<StaticImport> staticImports() {
-        return Collections.singleton(StaticImport.staticMethodImport(Collectors.class, "toSet"));
-    }
-
-    private MethodSpec createConstructor() {
-        return MethodSpec.constructorBuilder()
-                         .addModifiers(Modifier.PRIVATE)
-                         .addParameter(String.class, VALUE)
-                         .addStatement("this.$1N = $1N", VALUE)
-                         .build();
-    }
-
-    private MethodSpec fromValueSpec() {
-        return MethodSpec.methodBuilder("fromValue")
-                         .returns(className)
-                         .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
-                         .addJavadoc("Use this in place of valueOf to convert the raw string returned by the service into the " +
-                                     "enum value.\n\n" +
-                                     "@param $N real value\n" +
-                                     "@return $T corresponding to the value\n", VALUE, className)
-                         .addParameter(String.class, VALUE)
-                         .beginControlFlow("if ($N == null)", VALUE)
-                         .addStatement("return null")
-                         .endControlFlow()
-                         .addStatement("return $1T.of($2T.values())\n" +
-                                       ".filter(e -> e.toString().equals($3N))\n" +
-                                       ".findFirst()\n" +
-                                       ".orElse(UNKNOWN_TO_SDK_VERSION)",
-                                       Stream.class,
-                                       className,
-                                       VALUE)
-                         .build();
+    protected void addEnumConstants(Builder enumBuilder) {
+        getShape().getEnums().forEach(
+            e -> enumBuilder.addEnumConstant(e.getName(), TypeSpec.anonymousClassBuilder("$S", e.getValue()).build())
+        );
     }
 
-    private MethodSpec knownValuesSpec() {
-        return MethodSpec.methodBuilder("knownValues")
-                         .returns(ParameterizedTypeName.get(ClassName.get(Set.class), className))
-                         .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
-                         .addJavadoc("Use this in place of {@link #values()} to return a {@link Set} of all values known to the "
-                                     + "SDK.\n"
-                                     + "This will return all known enum values except {@link #$N}.\n\n"
-                                     + "@return a {@link $T} of known {@link $T}s", UNKNOWN_TO_SDK_VERSION, Set.class, className)
-                         .addStatement("return $T.of(values()).filter(v -> v != $N).collect(toSet())",
-                                       Stream.class,
-                                       UNKNOWN_TO_SDK_VERSION)
-                         .build();
+    @Override
+    public ClassName className() {
+        return ClassName.get(enumPackageName, getShape().getShapeName());
     }
 }

codegen/src/main/java/software/amazon/awssdk/codegen/poet/eventstream/EventStreamUtils.java

@@ -16,6 +16,7 @@
 package software.amazon.awssdk.codegen.poet.eventstream;
 
 import java.util.Collection;
+import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.stream.Collectors;
@@ -95,6 +96,14 @@ public static Optional<ShapeModel> getBaseEventStreamShape(IntermediateModel mod
                     .findFirst();
     }
 
+    public static List<ShapeModel> getBaseEventStreamShapes(IntermediateModel model, ShapeModel eventShape) {
+        return model.getShapes().values()
+                .stream()
+                .filter(ShapeModel::isEventStream)
+                .filter(s -> s.getMembers().stream().anyMatch(m -> m.getShape().equals(eventShape)))
+                .collect(Collectors.toList());
+    }
+
     /**
      * Returns the stream of event members ('event: true') excluding exceptions
      * from the input event stream shape ('eventstream: true').

codegen/src/main/java/software/amazon/awssdk/codegen/poet/eventstream/EventTypeEnumSpec.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.codegen.poet.eventstream;
+
+import com.squareup.javapoet.ClassName;
+import com.squareup.javapoet.TypeSpec;
+import software.amazon.awssdk.codegen.model.intermediate.IntermediateModel;
+import software.amazon.awssdk.codegen.model.intermediate.ShapeModel;
+import software.amazon.awssdk.codegen.naming.NamingStrategy;
+import software.amazon.awssdk.codegen.poet.common.AbstractEnumClass;
+
+public class EventTypeEnumSpec extends AbstractEnumClass {
+    private final String enumPackageName;
+    private final IntermediateModel intermediateModel;
+
+    public EventTypeEnumSpec(String enumPackageName, IntermediateModel intermediateModel, ShapeModel eventStream) {
+        super(eventStream);
+        this.enumPackageName = enumPackageName;
+        this.intermediateModel = intermediateModel;
+    }
+
+    @Override
+    protected void addDeprecated(TypeSpec.Builder enumBuilder) {
+        // no-op
+    }
+
+    @Override
+    protected void addJavadoc(TypeSpec.Builder enumBuilder) {
+        enumBuilder.addJavadoc("The known possible types of events for {@code $N}.", getShape().getShapeName());
+    }
+
+    @Override
+    protected void addEnumConstants(TypeSpec.Builder enumBuilder) {
+        NamingStrategy namingStrategy = intermediateModel.getNamingStrategy();
+        getShape().getMembers().stream()
+                .filter(m -> m.getShape().isEvent())
+                .forEach(m -> {
+                    String value = m.getC2jName();
+                    String name = namingStrategy.getEnumValueName(value);
+                    enumBuilder.addEnumConstant(name, TypeSpec.anonymousClassBuilder("$S", value).build());
+                });
+    }
+
+    @Override
+    public ClassName className() {
+        return ClassName.get(enumPackageName, "EventType");
+    }
+}