java generator, support non-string expandable enum for TypeSpec

packages/http-client-java/generator/http-client-generator-core/src/main/java/com/microsoft/typespec/http/client/generator/core/mapper/ChoiceMapper.java

@@ -4,6 +4,7 @@
 package com.microsoft.typespec.http.client.generator.core.mapper;
 
 import com.microsoft.typespec.http.client.generator.core.extension.model.codemodel.ChoiceSchema;
+import com.microsoft.typespec.http.client.generator.core.model.clientmodel.ClassType;
 import com.microsoft.typespec.http.client.generator.core.model.clientmodel.EnumType;
 import com.microsoft.typespec.http.client.generator.core.model.clientmodel.IType;
 import java.util.Map;
@@ -45,7 +46,17 @@ public IType map(ChoiceSchema enumType) {
         return choiceType;
     }
 
+    protected boolean useCodeModelNameForEnumMember() {
+        return true;
+    }
+
     private IType createChoiceType(ChoiceSchema enumType) {
-        return MapperUtils.createEnumType(enumType, true, true);
+        IType elementType = Mappers.getSchemaMapper().map(enumType.getChoiceType());
+        boolean isStringEnum = elementType == ClassType.STRING;
+        if (isStringEnum) {
+            return MapperUtils.createEnumType(enumType, true, useCodeModelNameForEnumMember());
+        } else {
+            return MapperUtils.createEnumType(enumType, true, useCodeModelNameForEnumMember(), "getValue", "fromValue");
+        }
     }
 }

packages/http-client-java/generator/http-client-generator-core/src/main/java/com/microsoft/typespec/http/client/generator/core/mapper/MapperUtils.java

@@ -29,14 +29,29 @@
 public final class MapperUtils {
     /**
      * Create enum client type from code model.
-     * 
+     *
      * @param enumType code model schema for enum
      * @param expandable whether it's expandable enum
      * @param useCodeModelNameForEnumMember whether to use code model enum member name for client enum member name
      * @return enum client type
      */
     public static IType createEnumType(ChoiceSchema enumType, boolean expandable,
         boolean useCodeModelNameForEnumMember) {
+        return createEnumType(enumType, expandable, useCodeModelNameForEnumMember, null, null);
+    }
+
+    /**
+     * Create enum client type from code model.
+     *
+     * @param enumType code model schema for enum
+     * @param expandable whether it's expandable enum
+     * @param useCodeModelNameForEnumMember whether to use code model enum member name for client enum member name
+     * @param serializationMethodName method name for serialization
+     * @param deserializationMethodName method name for deserialization
+     * @return enum client type
+     */
+    public static IType createEnumType(ChoiceSchema enumType, boolean expandable, boolean useCodeModelNameForEnumMember,
+        String serializationMethodName, String deserializationMethodName) {
         JavaSettings settings = JavaSettings.getInstance();
         String enumTypeName = enumType.getLanguage().getJava().getName();
 
@@ -98,6 +113,8 @@ public static IType createEnumType(ChoiceSchema enumType, boolean expandable,
                     new ImplementationDetails.Builder().usages(SchemaUtil.mapSchemaContext(enumType.getUsage()))
                         .build())
                 .crossLanguageDefinitionId(enumType.getCrossLanguageDefinitionId())
+                .fromMethodName(deserializationMethodName)
+                .toMethodName(serializationMethodName)
                 .build();
         }
     }

packages/http-client-java/generator/http-client-generator-core/src/main/java/com/microsoft/typespec/http/client/generator/core/mapper/SchemaMapper.java

@@ -49,10 +49,10 @@ public IType map(Schema value) {
     private IType createSchemaType(Schema value) {
         if (value instanceof PrimitiveSchema) {
             return Mappers.getPrimitiveMapper().map((PrimitiveSchema) value);
-        } else if (value instanceof ChoiceSchema) {
-            return Mappers.getChoiceMapper().map((ChoiceSchema) value);
         } else if (value instanceof SealedChoiceSchema) {
             return Mappers.getSealedChoiceMapper().map((SealedChoiceSchema) value);
+        } else if (value instanceof ChoiceSchema) {
+            return Mappers.getChoiceMapper().map((ChoiceSchema) value);
         } else if (value instanceof ArraySchema) {
             return Mappers.getArrayMapper().map((ArraySchema) value);
         } else if (value instanceof DictionarySchema) {

packages/http-client-java/generator/http-client-generator-core/src/main/java/com/microsoft/typespec/http/client/generator/core/mapper/SealedChoiceMapper.java

@@ -36,7 +36,11 @@ public IType map(SealedChoiceSchema enumType) {
         return sealedChoiceType;
     }
 
+    protected boolean useCodeModelNameForEnumMember() {
+        return true;
+    }
+
     private IType createSealedChoiceType(SealedChoiceSchema enumType) {
-        return MapperUtils.createEnumType(enumType, false, true);
+        return MapperUtils.createEnumType(enumType, false, useCodeModelNameForEnumMember());
     }
 }

packages/http-client-java/generator/http-client-generator-core/src/main/java/com/microsoft/typespec/http/client/generator/core/model/clientmodel/ClassType.java

@@ -45,6 +45,7 @@
 import com.azure.core.util.Context;
 import com.azure.core.util.CoreUtils;
 import com.azure.core.util.DateTimeRfc1123;
+import com.azure.core.util.ExpandableEnum;
 import com.azure.core.util.ExpandableStringEnum;
 import com.azure.core.util.logging.ClientLogger;
 import com.azure.core.util.logging.LogLevel;
@@ -178,6 +179,7 @@ private static ClassType.Builder getClassTypeBuilder(Class<?> classKey) {
     public static final ClassType RESPONSE = getClassTypeBuilder(Response.class).build();
     public static final ClassType SIMPLE_RESPONSE = getClassTypeBuilder(SimpleResponse.class).build();
     public static final ClassType EXPANDABLE_STRING_ENUM = getClassTypeBuilder(ExpandableStringEnum.class).build();
+    public static final ClassType EXPANDABLE_ENUM = getClassTypeBuilder(ExpandableEnum.class).build();
     public static final ClassType HTTP_PIPELINE_BUILDER = getClassTypeBuilder(HttpPipelineBuilder.class).build();
     public static final ClassType KEY_CREDENTIAL_POLICY = getClassTypeBuilder(KeyCredentialPolicy.class).build();
     public static final ClassType KEY_CREDENTIAL_TRAIT = getClassTypeBuilder(KeyCredentialTrait.class).build();

packages/http-client-java/generator/http-client-generator-core/src/main/java/com/microsoft/typespec/http/client/generator/core/model/clientmodel/EnumType.java

@@ -3,6 +3,7 @@
 
 package com.microsoft.typespec.http.client.generator.core.model.clientmodel;
 
+import com.azure.core.util.CoreUtils;
 import com.microsoft.typespec.http.client.generator.core.util.CodeNamer;
 import java.util.List;
 import java.util.Set;
@@ -34,18 +35,22 @@ public class EnumType implements IType {
     private final ImplementationDetails implementationDetails;
 
     private String crossLanguageDefinitionId;
+    private final String fromMethodName;
+    private final String toMethodName;
 
     /**
      * Create a new Enum with the provided properties.
-     * 
+     *
      * @param name The name of the new Enum.
      * @param description The description of the Enum.
      * @param expandable Whether this will be an ExpandableStringEnum type.
      * @param values The values of the Enum.
+     * @param fromMethodName The method name used to convert JSON to the enum type.
+     * @param toMethodName The method name used to convert the enum type to JSON.
      */
     private EnumType(String packageKeyword, String name, String description, boolean expandable,
         List<ClientEnumValue> values, IType elementType, ImplementationDetails implementationDetails,
-        String crossLanguageDefinitionId) {
+        String crossLanguageDefinitionId, String fromMethodName, String toMethodName) {
         this.name = name;
         this.packageName = packageKeyword;
         this.description = description;
@@ -54,6 +59,8 @@ private EnumType(String packageKeyword, String name, String description, boolean
         this.elementType = elementType;
         this.implementationDetails = implementationDetails;
         this.crossLanguageDefinitionId = crossLanguageDefinitionId;
+        this.fromMethodName = fromMethodName;
+        this.toMethodName = toMethodName;
     }
 
     public String getCrossLanguageDefinitionId() {
@@ -132,7 +139,9 @@ public final String defaultValueExpression(String sourceExpression) {
      * @return The method name used to convert JSON to the enum type.
      */
     public final String getFromMethodName() {
-        return "from" + CodeNamer.toPascalCase(elementType.getClientType().toString());
+        return CoreUtils.isNullOrEmpty(fromMethodName)
+            ? "from" + CodeNamer.toPascalCase(elementType.getClientType().toString())
+            : fromMethodName;
     }
 
     /**
@@ -141,7 +150,9 @@ public final String getFromMethodName() {
      * @return The method name used to convert the enum type to JSON.
      */
     public final String getToMethodName() {
-        return "to" + CodeNamer.toPascalCase(elementType.getClientType().toString());
+        return CoreUtils.isNullOrEmpty(toMethodName)
+            ? "to" + CodeNamer.toPascalCase(elementType.getClientType().toString())
+            : toMethodName;
     }
 
     @Override
@@ -229,6 +240,8 @@ public static class Builder {
         private ImplementationDetails implementationDetails;
 
         private String crossLanguageDefinitionId;
+        private String fromMethodName;
+        private String toMethodName;
 
         /**
          * Sets the name of the Enum.
@@ -314,12 +327,22 @@ public Builder crossLanguageDefinitionId(String crossLanguageDefinitionId) {
             return this;
         }
 
+        public Builder fromMethodName(String fromMethodName) {
+            this.fromMethodName = fromMethodName;
+            return this;
+        }
+
+        public Builder toMethodName(String toMethodName) {
+            this.toMethodName = toMethodName;
+            return this;
+        }
+
         /**
          * @return an immutable EnumType instance with the configurations on this builder.
          */
         public EnumType build() {
             return new EnumType(packageName, name, description, expandable, values, elementType, implementationDetails,
-                crossLanguageDefinitionId);
+                crossLanguageDefinitionId, fromMethodName, toMethodName);
         }
     }
 }

packages/http-client-java/generator/http-client-generator-core/src/main/java/com/microsoft/typespec/http/client/generator/core/template/EnumTemplate.java

@@ -39,7 +39,7 @@ public final void write(EnumType enumType, JavaFile javaFile) {
 
         if (enumType.getExpandable()) {
             if (settings.isBranded()) {
-                writeExpandableStringEnum(enumType, javaFile, settings);
+                writeBrandedExpandableEnum(enumType, javaFile, settings);
             } else {
                 writeExpandableStringEnumInterface(enumType, javaFile, settings);
             }
@@ -48,6 +48,119 @@ public final void write(EnumType enumType, JavaFile javaFile) {
         }
     }
 
+    /**
+     * Extension point for expandable enum implementation of branded flavor.
+     *
+     * @param enumType enumType to write implementation
+     * @param javaFile javaFile to write into
+     * @param settings {@link JavaSettings} instance
+     */
+    protected void writeBrandedExpandableEnum(EnumType enumType, JavaFile javaFile, JavaSettings settings) {
+        if (enumType.getElementType() == ClassType.STRING) {
+            writeExpandableStringEnum(enumType, javaFile, settings);
+        } else {
+            Set<String> imports = new HashSet<>();
+            imports.add("java.util.Collection");
+            imports.add("java.lang.IllegalArgumentException");
+            imports.add("java.util.Map");
+            imports.add("java.util.concurrent.ConcurrentHashMap");
+            imports.add("java.util.ArrayList");
+            imports.add("java.util.Objects");
+            imports.add(ClassType.EXPANDABLE_ENUM.getFullName());
+            if (!settings.isStreamStyleSerialization()) {
+                imports.add("com.fasterxml.jackson.annotation.JsonCreator");
+            }
+
+            addGeneratedImport(imports);
+
+            javaFile.declareImport(imports);
+            javaFile.javadocComment(comment -> comment.description(enumType.getDescription()));
+
+            String enumName = enumType.getName();
+            IType elementType = enumType.getElementType();
+            String typeName = elementType.getClientType().asNullable().toString();
+            String pascalTypeName = CodeNamer.toPascalCase(typeName);
+            String declaration = enumName + " implements ExpandableEnum<" + pascalTypeName + ">";
+            javaFile.publicFinalClass(declaration, classBlock -> {
+                classBlock.privateStaticFinalVariable(
+                    String.format("Map<%1$s, %2$s> VALUES = new ConcurrentHashMap<>()", pascalTypeName, enumName));
+
+                for (ClientEnumValue enumValue : enumType.getValues()) {
+                    String value = enumValue.getValue();
+                    classBlock.javadocComment(CoreUtils.isNullOrEmpty(enumValue.getDescription())
+                        ? "Static value " + value + " for " + enumName + "."
+                        : enumValue.getDescription());
+                    addGeneratedAnnotation(classBlock);
+                    classBlock.publicStaticFinalVariable(String.format("%1$s %2$s = fromValue(%3$s)", enumName,
+                        enumValue.getName(), elementType.defaultValueExpression(value)));
+                }
+
+                classBlock.variable(pascalTypeName + " value", JavaVisibility.Private, JavaModifier.Final);
+                classBlock.privateConstructor(enumName + "(" + pascalTypeName + " value)", ctor -> {
+                    ctor.line("this.value = value;");
+                });
+
+                // fromValue(typeName)
+                classBlock.javadocComment(comment -> {
+                    comment.description("Creates or finds a " + enumName);
+                    comment.param("value", "a value to look for");
+                    comment.methodReturns("the corresponding " + enumName);
+                });
+
+                addGeneratedAnnotation(classBlock);
+                if (!settings.isStreamStyleSerialization()) {
+                    classBlock.annotation("JsonCreator");
+                }
+
+                classBlock.publicStaticMethod(String.format("%1$s fromValue(%2$s value)", enumName, pascalTypeName),
+                    function -> {
+                        function.line("Objects.requireNonNull(value, \"'value' cannot be null.\");");
+                        function.line(enumName + " member = VALUES.get(value);");
+                        function.ifBlock("member != null", ifAction -> ifAction.line("return member;"));
+                        function.methodReturn("VALUES.computeIfAbsent(value, key -> new " + enumName + "(key))");
+                    });
+
+                // values
+                classBlock.javadocComment(comment -> {
+                    comment.description("Gets known " + enumName + " values.");
+                    comment.methodReturns("Known " + enumName + " values.");
+                });
+                addGeneratedAnnotation(classBlock);
+                classBlock.publicStaticMethod(String.format("Collection<%s> values()", enumName),
+                    function -> function.methodReturn("new ArrayList<>(VALUES.values())"));
+
+                // getValue
+                classBlock.javadocComment(comment -> {
+                    comment.description("Gets the value of the " + enumName + " instance.");
+                    comment.methodReturns("the value of the " + enumName + " instance.");
+                });
+
+                addGeneratedAnnotation(classBlock);
+                classBlock.annotation("Override");
+                classBlock.publicMethod(pascalTypeName + " getValue()",
+                    function -> function.methodReturn("this.value"));
+
+                // toString
+                addGeneratedAnnotation(classBlock);
+                classBlock.annotation("Override");
+                classBlock.method(JavaVisibility.Public, null, "String toString()",
+                    function -> function.methodReturn("Objects.toString(this.value)"));
+
+                // equals
+                addGeneratedAnnotation(classBlock);
+                classBlock.annotation("Override");
+                classBlock.method(JavaVisibility.Public, null, "boolean equals(Object obj)",
+                    function -> function.methodReturn("Objects.equals(this.value, obj)"));
+
+                // hashcode
+                addGeneratedAnnotation(classBlock);
+                classBlock.annotation("Override");
+                classBlock.method(JavaVisibility.Public, null, "int hashCode()",
+                    function -> function.methodReturn("Objects.hashCode(this.value)"));
+            });
+        }
+    }
+
     private void writeExpandableStringEnumInterface(EnumType enumType, JavaFile javaFile, JavaSettings settings) {
         Set<String> imports = new HashSet<>();
         imports.add("java.util.Collection");

packages/http-client-java/generator/http-client-generator-mgmt/src/main/java/com/microsoft/typespec/http/client/generator/mgmt/mapper/FluentChoiceMapper.java

@@ -3,10 +3,7 @@
 
 package com.microsoft.typespec.http.client.generator.mgmt.mapper;
 
-import com.microsoft.typespec.http.client.generator.core.extension.model.codemodel.ChoiceSchema;
 import com.microsoft.typespec.http.client.generator.core.mapper.ChoiceMapper;
-import com.microsoft.typespec.http.client.generator.core.mapper.MapperUtils;
-import com.microsoft.typespec.http.client.generator.core.model.clientmodel.IType;
 
 public class FluentChoiceMapper extends ChoiceMapper {
     private static final FluentChoiceMapper INSTANCE = new FluentChoiceMapper();
@@ -19,7 +16,7 @@ public static FluentChoiceMapper getInstance() {
     }
 
     @Override
-    public IType map(ChoiceSchema enumType) {
-        return MapperUtils.createEnumType(enumType, true, false);
+    protected boolean useCodeModelNameForEnumMember() {
+        return false;
     }
 }

packages/http-client-java/generator/http-client-generator-mgmt/src/main/java/com/microsoft/typespec/http/client/generator/mgmt/mapper/FluentSealedChoiceMapper.java

@@ -3,10 +3,7 @@
 
 package com.microsoft.typespec.http.client.generator.mgmt.mapper;
 
-import com.microsoft.typespec.http.client.generator.core.extension.model.codemodel.SealedChoiceSchema;
-import com.microsoft.typespec.http.client.generator.core.mapper.MapperUtils;
 import com.microsoft.typespec.http.client.generator.core.mapper.SealedChoiceMapper;
-import com.microsoft.typespec.http.client.generator.core.model.clientmodel.IType;
 
 public class FluentSealedChoiceMapper extends SealedChoiceMapper {
     private static final FluentSealedChoiceMapper INSTANCE = new FluentSealedChoiceMapper();
@@ -19,7 +16,7 @@ public static FluentSealedChoiceMapper getInstance() {
     }
 
     @Override
-    public IType map(SealedChoiceSchema enumType) {
-        return MapperUtils.createEnumType(enumType, false, false);
+    protected boolean useCodeModelNameForEnumMember() {
+        return false;
     }
 }