Adding knownTypes static to enum generation to filter out the unknown SDK enums.

kiiadi · kiiadi · commit ba0017c0daa3 · 2017-10-18T12:36:49.000-07:00
diff --git a/codegen/src/main/java/software/amazon/awssdk/codegen/poet/common/EnumClass.java b/codegen/src/main/java/software/amazon/awssdk/codegen/poet/common/EnumClass.java
@@ -22,16 +22,22 @@
 
 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 {
 
     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;
 
@@ -46,6 +52,7 @@ public TypeSpec poetSpec() {
             .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);
@@ -54,7 +61,7 @@ public TypeSpec poetSpec() {
         shape.getEnums().forEach(
             e -> enumBuilder.addEnumConstant(e.getName(), TypeSpec.anonymousClassBuilder("$S", e.getValue()).build())
         );
-        enumBuilder.addEnumConstant("UNKNOWN_TO_SDK_VERSION", TypeSpec.anonymousClassBuilder("null").build());
+        enumBuilder.addEnumConstant(UNKNOWN_TO_SDK_VERSION, TypeSpec.anonymousClassBuilder("null").build());
 
         return enumBuilder.build();
     }
@@ -64,6 +71,11 @@ public ClassName className() {
         return className;
     }
 
+    @Override
+    public Iterable<StaticImport> staticImports() {
+        return Collections.singleton(StaticImport.staticMethodImport(Collectors.class, "toSet"));
+    }
+
     private MethodSpec createConstructor() {
         return MethodSpec.constructorBuilder()
                          .addModifiers(Modifier.PRIVATE)
@@ -93,4 +105,18 @@ private MethodSpec fromValueSpec() {
                                        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(toSet())",
+                                       Stream.class,
+                                       UNKNOWN_TO_SDK_VERSION)
+                         .build();
+    }
 }
diff --git a/codegen/src/main/resources/templates/common/enum-class.ftl b/codegen/src/main/resources/templates/common/enum-class.ftl
@@ -49,4 +49,8 @@ public enum ${shape.shapeName} {
 
         throw new IllegalArgumentException("Cannot create enum from " + value + " value!");
     }
+
+    public static Set<${shape.shapeName}> knownValues() {
+        return Stream.of(values()).filter(v -> v != UNKNOWN_TO_SDK_VERSION).collect(toSet());
+    }
 }
diff --git a/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/common/test-enum-class.java b/codegen/src/test/resources/software/amazon/awssdk/codegen/poet/common/test-enum-class.java
@@ -1,5 +1,8 @@
 package software.amazon.awssdk.codegen.poet.common.model;
 
+import static java.util.stream.Collectors.toSet;
+
+import java.util.Set;
 import java.util.stream.Stream;
 import javax.annotation.Generated;
 
@@ -38,4 +41,14 @@ public static TestEnumClass fromValue(String value) {
         }
         return Stream.of(TestEnumClass.values()).filter(e -> e.toString().equals(value)).findFirst().orElse(UNKNOWN_TO_SDK_VERSION);
     }
+
+    /**
+     * Use this in place of {@link #values()} to return a {@link Set} of all values known to the SDK.
+     * This will return all known enum values except {@link #UNKNOWN_TO_SDK_VERSION}.
+     *
+     * @return a {@link Set} of known {@link TestEnumClass}s
+     */
+    public static Set<TestEnumClass> knownValues() {
+        return Stream.of(values()).filter(v -> v != UNKNOWN_TO_SDK_VERSION).collect(toSet());
+    }
 }

Original file line number	Diff line number	Diff line change
`@@ -49,4 +49,8 @@ public enum ${shape.shapeName} {`
`49`	`49`
`50`	`50`	`throw new IllegalArgumentException("Cannot create enum from " + value + " value!");`
`51`	`51`	`}`
	`52`	`+`
	`53`	`+ public static Set<${shape.shapeName}> knownValues() {`
	`54`	`+ return Stream.of(values()).filter(v -> v != UNKNOWN_TO_SDK_VERSION).collect(toSet());`
	`55`	`+ }`
`52`	`56`	`}`