Generate visitor interfaces for conjure enums (#142)

Most uses of enums involve a switch or if-else block across values along the lines of "boolean isTerminalState(enum)" or "String getDescription(enum)". It's easy to add new enum values without updating all usages, by using a visitor we catch these bugs at compile time rather than runtime.
palantir · Jan 9, 2019 · 5994e94 · 5994e94
1 parent 8e4f032
commit 5994e94
Show file tree

Hide file tree

Showing 4 changed files with 181 additions and 7 deletions.
diff --git a/conjure-java-core/src/integrationInput/java/com/palantir/product/EnumExample.java b/conjure-java-core/src/integrationInput/java/com/palantir/product/EnumExample.java
@@ -7,7 +7,7 @@
 import javax.annotation.Generated;
 
 /**
- * This enumerates the numbers 1:2.
+ * This enumerates the numbers 1:2 also 100.
  *
  * <p>This class is used instead of a native enum to support unknown values. Rather than throw an
  * exception, the {@link EnumExample#valueOf} method defaults to a new instantiation of {@link
@@ -25,6 +25,9 @@ public final class EnumExample {
 
     public static final EnumExample TWO = new EnumExample(Value.TWO, "TWO");
 
+    /** Value of 100. */
+    public static final EnumExample ONE_HUNDRED = new EnumExample(Value.ONE_HUNDRED, "ONE_HUNDRED");
+
     private final Value value;
 
     private final String string;
@@ -65,17 +68,47 @@ public static EnumExample valueOf(String value) {
                 return ONE;
             case "TWO":
                 return TWO;
+            case "ONE_HUNDRED":
+                return ONE_HUNDRED;
             default:
                 return new EnumExample(Value.UNKNOWN, upperCasedValue);
         }
     }
 
+    public <T> T accept(Visitor<T> visitor) {
+        switch (value) {
+            case ONE:
+                return visitor.visitOne();
+            case TWO:
+                return visitor.visitTwo();
+            case ONE_HUNDRED:
+                return visitor.visitOneHundred();
+            default:
+                return visitor.visitUnknown(string);
+        }
+    }
+
     @Generated("com.palantir.conjure.java.types.EnumGenerator")
     public enum Value {
         ONE,
 
         TWO,
 
+        /** Value of 100. */
+        ONE_HUNDRED,
+
         UNKNOWN
     }
+
+    @Generated("com.palantir.conjure.java.types.EnumGenerator")
+    public interface Visitor<T> {
+        T visitOne();
+
+        T visitTwo();
+
+        /** Value of 100. */
+        T visitOneHundred();
+
+        T visitUnknown(String unknownValue);
+    }
 }
diff --git a/conjure-java-core/src/main/java/com/palantir/conjure/java/types/EnumGenerator.java b/conjure-java-core/src/main/java/com/palantir/conjure/java/types/EnumGenerator.java
@@ -18,6 +18,8 @@
 
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonValue;
+import com.google.common.base.CaseFormat;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.palantir.conjure.java.ConjureAnnotations;
 import com.palantir.conjure.spec.EnumDefinition;
@@ -28,34 +30,46 @@
 import com.squareup.javapoet.JavaFile;
 import com.squareup.javapoet.MethodSpec;
 import com.squareup.javapoet.ParameterSpec;
+import com.squareup.javapoet.ParameterizedTypeName;
 import com.squareup.javapoet.TypeName;
 import com.squareup.javapoet.TypeSpec;
+import com.squareup.javapoet.TypeVariableName;
+import java.util.List;
 import java.util.Locale;
 import javax.lang.model.element.Modifier;
 import org.apache.commons.lang3.StringUtils;
 
 public final class EnumGenerator {
 
+    private static final String VALUE_PARAMETER = "value";
+    private static final String STRING_PARAMETER = "string";
+    private static final String VISIT_METHOD_NAME = "visit";
+    private static final String VISIT_UNKNOWN_METHOD_NAME = "visitUnknown";
+    private static final TypeVariableName TYPE_VARIABLE = TypeVariableName.get("T");
+
     private EnumGenerator() {}
 
     public static JavaFile generateEnumType(EnumDefinition typeDef) {
         String typePackage = typeDef.getTypeName().getPackage();
         ClassName thisClass = ClassName.get(typePackage, typeDef.getTypeName().getName());
         ClassName enumClass = ClassName.get(typePackage, typeDef.getTypeName().getName(), "Value");
+        ClassName visitorClass = ClassName.get(typePackage, typeDef.getTypeName().getName(), "Visitor");
 
-        return JavaFile.builder(typePackage, createSafeEnum(typeDef, thisClass, enumClass))
+        return JavaFile.builder(typePackage, createSafeEnum(typeDef, thisClass, enumClass, visitorClass))
                 .skipJavaLangImports(true)
                 .indent("    ")
                 .build();
     }
 
-    private static TypeSpec createSafeEnum(EnumDefinition typeDef, ClassName thisClass, ClassName enumClass) {
+    private static TypeSpec createSafeEnum(
+            EnumDefinition typeDef, ClassName thisClass, ClassName enumClass, ClassName visitorClass) {
         TypeSpec.Builder wrapper = TypeSpec.classBuilder(typeDef.getTypeName().getName())
                 .addAnnotation(ConjureAnnotations.getConjureGeneratedAnnotation(EnumGenerator.class))
                 .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                 .addType(createEnum(enumClass, typeDef.getValues(), true))
-                .addField(enumClass, "value", Modifier.PRIVATE, Modifier.FINAL)
-                .addField(ClassName.get(String.class), "string", Modifier.PRIVATE, Modifier.FINAL)
+                .addType(createVisitor(visitorClass, typeDef.getValues()))
+                .addField(enumClass, VALUE_PARAMETER, Modifier.PRIVATE, Modifier.FINAL)
+                .addField(ClassName.get(String.class), STRING_PARAMETER, Modifier.PRIVATE, Modifier.FINAL)
                 .addFields(createConstants(typeDef.getValues(), thisClass, enumClass))
                 .addMethod(createConstructor(enumClass))
                 .addMethod(MethodSpec.methodBuilder("get")
@@ -72,7 +86,8 @@ private static TypeSpec createSafeEnum(EnumDefinition typeDef, ClassName thisCla
                         .build())
                 .addMethod(createEquals(thisClass))
                 .addMethod(createHashCode())
-                .addMethod(createValueOf(thisClass, typeDef.getValues()));
+                .addMethod(createValueOf(thisClass, typeDef.getValues()))
+                .addMethod(generateAcceptVisitMethod(visitorClass, typeDef.getValues()));
 
         typeDef.getDocs().ifPresent(
                 docs -> wrapper.addJavadoc("$L<p>\n", StringUtils.appendIfMissing(docs.get(), "\n")));
@@ -133,6 +148,57 @@ private static TypeSpec createEnum(ClassName enumClass, Iterable<EnumValueDefini
         return enumBuilder.build();
     }
 
+    private static TypeSpec createVisitor(ClassName visitorClass, Iterable<EnumValueDefinition> values) {
+        return TypeSpec.interfaceBuilder(visitorClass.simpleName())
+                .addAnnotation(ConjureAnnotations.getConjureGeneratedAnnotation(EnumGenerator.class))
+                .addModifiers(Modifier.PUBLIC)
+                .addTypeVariable(TYPE_VARIABLE)
+                .addMethods(generateMemberVisitMethods(values))
+                .addMethod(MethodSpec.methodBuilder(VISIT_UNKNOWN_METHOD_NAME)
+                        .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)
+                        .addParameter(String.class, "unknownValue")
+                        .returns(TYPE_VARIABLE)
+                        .build())
+                .build();
+    }
+
+    private static List<MethodSpec> generateMemberVisitMethods(Iterable<EnumValueDefinition> values) {
+        ImmutableList.Builder<MethodSpec> methods = ImmutableList.builder();
+        for (EnumValueDefinition value : values) {
+            MethodSpec.Builder methodSpecBuilder = MethodSpec.methodBuilder(getVisitorMethodName(value))
+                    .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)
+                    .returns(TYPE_VARIABLE);
+            value.getDocs().ifPresent(docs ->
+                    methodSpecBuilder.addJavadoc("$L", StringUtils.appendIfMissing(docs.get(), "\n")));
+            methods.add(methodSpecBuilder.build());
+        }
+        return methods.build();
+    }
+
+    private static MethodSpec generateAcceptVisitMethod(ClassName visitorClass, Iterable<EnumValueDefinition> values) {
+        CodeBlock.Builder switchBlock = CodeBlock.builder();
+        switchBlock.beginControlFlow("switch (value)");
+        for (EnumValueDefinition value : values) {
+            switchBlock.add("case $N:\n", value.getValue())
+                    .addStatement("return visitor.$L()", getVisitorMethodName(value));
+        }
+        switchBlock.add("default:\n")
+                .addStatement("return visitor.$1L(string)", VISIT_UNKNOWN_METHOD_NAME)
+                .endControlFlow();
+        ParameterizedTypeName parameterizedVisitorClass = ParameterizedTypeName.get(visitorClass, TYPE_VARIABLE);
+        return MethodSpec.methodBuilder("accept")
+                .addModifiers(Modifier.PUBLIC)
+                .addParameter(ParameterSpec.builder(parameterizedVisitorClass, "visitor").build())
+                .addTypeVariable(TYPE_VARIABLE)
+                .returns(TYPE_VARIABLE)
+                .addCode(switchBlock.build())
+                .build();
+    }
+
+    private static String getVisitorMethodName(EnumValueDefinition definition) {
+        return VISIT_METHOD_NAME + CaseFormat.UPPER_UNDERSCORE.to(CaseFormat.UPPER_CAMEL, definition.getValue());
+    }
+
     private static MethodSpec createConstructor(ClassName enumClass) {
         // Note: We generate a two arg constructor that handles both known
         // and unknown variants instead of using separate contructors to avoid

diff --git a/conjure-java-core/src/test/java/com/palantir/conjure/java/types/EnumTests.java b/conjure-java-core/src/test/java/com/palantir/conjure/java/types/EnumTests.java
@@ -0,0 +1,73 @@
+/*
+ * (c) Copyright 2018 Palantir Technologies Inc. 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.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License 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 com.palantir.conjure.java.types;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import com.palantir.product.EnumExample;
+import org.junit.Test;
+
+public class EnumTests {
+
+    @Test
+    public void testVisitOne() {
+        EnumExample enumExample = EnumExample.ONE;
+        assertThat(enumExample.accept(Visitor.INSTANCE)).isEqualTo("one");
+    }
+
+    @Test
+    public void testValueOfProducesSameInstance() {
+        assertThat(EnumExample.valueOf("ONE")).isSameAs(EnumExample.ONE);
+    }
+
+    @Test
+    public void testUnknown() {
+        EnumExample enumExample = EnumExample.valueOf("SOME_VALUE");
+        assertThat(enumExample.get()).isEqualTo(EnumExample.Value.UNKNOWN);
+        assertThat(enumExample.toString()).isEqualTo("SOME_VALUE");
+    }
+
+    @Test
+    public void testVisitUnknown() {
+        EnumExample enumExample = EnumExample.valueOf("SOME_VALUE");
+        assertThat(enumExample.accept(Visitor.INSTANCE)).isEqualTo("SOME_VALUE");
+    }
+
+    private enum Visitor implements EnumExample.Visitor<String> {
+        INSTANCE;
+
+        @Override
+        public String visitOne() {
+            return "one";
+        }
+
+        @Override
+        public String visitTwo() {
+            return "two";
+        }
+
+        @Override
+        public String visitOneHundred() {
+            return "one hundred";
+        }
+
+        @Override
+        public String visitUnknown(String unknownValue) {
+            return unknownValue;
+        }
+    }
+}
diff --git a/conjure-java-core/src/test/resources/example-types.yml b/conjure-java-core/src/test/resources/example-types.yml
@@ -56,10 +56,12 @@ types:
           items: map<string, string>
       EnumExample:
         docs: |
-          This enumerates the numbers 1:2.
+          This enumerates the numbers 1:2 also 100.
         values:
           - ONE
           - TWO
+          - value: ONE_HUNDRED
+            docs: Value of 100.
       EnumFieldExample:
         fields:
           enum: EnumExample