OpenAPITools · johnlinp · Nov 27, 2021
diff --git a/modules/openapi-generator/src/main/java/org/openapitools/codegen/CodegenDiscriminator.java b/modules/openapi-generator/src/main/java/org/openapitools/codegen/CodegenDiscriminator.java
@@ -21,6 +21,8 @@ public class CodegenDiscriminator {
     private String propertyBaseName;
     private String propertyGetter;
     private String propertyType;
+    private boolean isEnumProperty;
+    private String mappingKey;
     private Map<String, String> mapping;
 
     // mappedModels is populated differently if legacyDiscriminatorBehavior is
@@ -72,6 +74,26 @@ public void setPropertyType(String propertyType) {
         this.propertyType = propertyType;
     }
 
+    public boolean getIsEnumProperty ()
+    {
+        return isEnumProperty;
+    }
+
+    public void setIsEnumProperty (boolean enumProperty)
+    {
+        isEnumProperty = enumProperty;
+    }
+
+    public String getMappingKey ()
+    {
+        return mappingKey;
+    }
+
+    public void setMappingKey (String mappingKey)
+    {
+        this.mappingKey = mappingKey;
+    }
+
     public Map<String, String> getMapping() {
         return mapping;
     }
@@ -180,4 +202,4 @@ public String toString() {
         sb.append('}');
         return sb.toString();
     }
-}
+}
diff --git a/modules/openapi-generator/src/main/java/org/openapitools/codegen/DefaultCodegen.java b/modules/openapi-generator/src/main/java/org/openapitools/codegen/DefaultCodegen.java
@@ -1501,7 +1501,18 @@ public String toArrayModelParamName(String name) {
      */
     @SuppressWarnings("static-method")
     public String toEnumName(CodegenProperty property) {
-        return StringUtils.capitalize(property.name) + "Enum";
+        return toEnumName(property.name);
+    }
+
+    /**
+     * Return the Enum name (e.g. StatusEnum given 'status')
+     *
+     * @param propertyName Property name
+     * @return the Enum name
+     */
+    @SuppressWarnings("static-method")
+    public String toEnumName(String propertyName) {
+        return StringUtils.capitalize(propertyName) + "Enum";
     }
 
     /**
@@ -3190,8 +3201,15 @@ protected CodegenDiscriminator createDiscriminator(String schemaName, Schema sch
         discriminator.setPropertyName(toVarName(discPropName));
         discriminator.setPropertyBaseName(sourceDiscriminator.getPropertyName());
         discriminator.setPropertyGetter(toGetter(discriminator.getPropertyName()));
-        // FIXME: for now, we assume that the discriminator property is String
-        discriminator.setPropertyType(typeMapping.get("string"));
+        if (isEnumDiscriminator(schema, openAPI)) {
+            discriminator.setPropertyType(toEnumName(discPropName));
+            discriminator.setIsEnumProperty(true);
+        } else {
+            // FIXME: let's assume that it's a string if it's not an enum
+            discriminator.setPropertyType(typeMapping.get("string"));
+            discriminator.setIsEnumProperty(false);
+        }
+        discriminator.setMappingKey(getMappingKey(schemaName, sourceDiscriminator.getMapping()));
         discriminator.setMapping(sourceDiscriminator.getMapping());
         List<MappedModel> uniqueDescendants = new ArrayList();
         if (sourceDiscriminator.getMapping() != null && !sourceDiscriminator.getMapping().isEmpty()) {
@@ -3245,6 +3263,50 @@ protected CodegenDiscriminator createDiscriminator(String schemaName, Schema sch
         return discriminator;
     }
 
+    private boolean isEnumDiscriminator (Schema schema, OpenAPI openAPI)
+    {
+        Schema refSchema = ModelUtils.getReferencedSchema(openAPI, schema);
+        Discriminator discriminator = refSchema.getDiscriminator();
+        if (discriminator != null) {
+            Map<String, Schema> properties = refSchema.getProperties();
+            if (properties != null) {
+                Schema discriminatorSchema = properties.get(discriminator.getPropertyName());
+                return discriminatorSchema != null &&
+                        discriminatorSchema.getEnum() != null
+                        && !discriminatorSchema.getEnum().isEmpty();
+            }
+        }
+
+        if (ModelUtils.isComposedSchema(refSchema)) {
+            ComposedSchema composedSchema = (ComposedSchema) refSchema;
+            if (composedSchema.getAllOf() != null) {
+                for (Schema allOf : composedSchema.getAllOf()) {
+                    if (isEnumDiscriminator(allOf, openAPI)) {
+                        return true;
+                    }
+                }
+            }
+        }
+
+        return false;
+    }
+
+    private String getMappingKey (String schemaName, Map<String, String> mapping)
+    {
+        if (mapping == null) {
+            return null;
+        }
+
+        for (Map.Entry<String, String> entry : mapping.entrySet()) {
+            String mappingSchemaPath = entry.getValue();
+            String mappingSchemaName = mappingSchemaPath.substring(mappingSchemaPath.lastIndexOf("/") + 1);
+            if (schemaName.equals(mappingSchemaName)) {
+                return entry.getKey();
+            }
+        }
+        return null;
+    }
+
     /**
      * Handle the model for the 'additionalProperties' keyword in the OAS schema.
      *

diff --git a/modules/openapi-generator/src/main/resources/Java/pojo.mustache b/modules/openapi-generator/src/main/resources/Java/pojo.mustache
@@ -79,7 +79,9 @@ public class {{classname}} {{#parent}}extends {{{.}}} {{/parent}}{{#vendorExtens
   {{/vars}}
   public {{classname}}() { {{#parent}}{{#parcelableModel}}
     super();{{/parcelableModel}}{{/parent}}{{#gson}}{{#discriminator}}
-    this.{{{discriminatorName}}} = this.getClass().getSimpleName();{{/discriminator}}{{/gson}}
+    {{#isEnumProperty}}{{#mappingKey}}this.{{{discriminatorName}}} = {{propertyType}}.fromValue("{{mappingKey}}");{{/mappingKey}}{{/isEnumProperty}}
+    {{^isEnumProperty}}this.{{{discriminatorName}}} = this.getClass().getSimpleName();{{/isEnumProperty}}
+    {{/discriminator}}{{/gson}}
   }{{#vendorExtensions.x-has-readonly-properties}}{{^withXml}}
 
   {{#jsonb}}@JsonbCreator{{/jsonb}}{{#jackson}}@JsonCreator{{/jackson}}

diff --git a/.../openapi-generator/src/test/java/org/openapitools/codegen/java/JavaClientCodegenTest.java b/.../openapi-generator/src/test/java/org/openapitools/codegen/java/JavaClientCodegenTest.java
@@ -1288,4 +1288,52 @@ public void testNativeClientWhiteSpacePathParamEncoding() throws IOException {
         TestUtils.assertFileContains(Paths.get(output + "/src/main/java/xyz/abcdef/ApiClient.java"),
                 "public static String urlEncode(String s) { return URLEncoder.encode(s, UTF_8).replaceAll(\"\\\\+\", \"%20\"); }");
     }
+
+    @Test
+    public void testStringDiscriminator() throws IOException {
+        File output = Files.createTempDirectory("test").toFile();
+
+        final CodegenConfigurator configurator = new CodegenConfigurator()
+                .setGeneratorName("java")
+                .setInputSpec("src/test/resources/3_0/issue_806_string_discriminator.yaml")
+                .setOutputDir(output.getAbsolutePath().replace("\\", "/"));
+
+        final ClientOptInput clientOptInput = configurator.toClientOptInput();
+        DefaultGenerator generator = new DefaultGenerator();
+        List<File> files = generator.opts(clientOptInput).generate();
+
+        validateJavaSourceFiles(files);
+
+        TestUtils.assertFileContains(Paths.get(output + "/src/main/java/org/openapitools/client/model/Dog.java"),
+                                     "this.petType = this.getClass().getSimpleName();");
+        TestUtils.assertFileContains(Paths.get(output + "/src/main/java/org/openapitools/client/model/Cat.java"),
+                                     "this.petType = this.getClass().getSimpleName();");
+        TestUtils.assertFileContains(Paths.get(output + "/src/main/java/org/openapitools/client/model/Lizard.java"),
+                                     "this.petType = this.getClass().getSimpleName();");
+
+        output.deleteOnExit();
+    }
+
+    @Test
+    public void testEnumDiscriminator() throws IOException {
+        File output = Files.createTempDirectory("test").toFile();
+
+        final CodegenConfigurator configurator = new CodegenConfigurator()
+                .setGeneratorName("java")
+                .setInputSpec("src/test/resources/3_0/issue_806_enum_discriminator.yaml")
+                .setOutputDir(output.getAbsolutePath().replace("\\", "/"));
+
+        final ClientOptInput clientOptInput = configurator.toClientOptInput();
+        DefaultGenerator generator = new DefaultGenerator();
+        List<File> files = generator.opts(clientOptInput).generate();
+
+        validateJavaSourceFiles(files);
+
+        TestUtils.assertFileContains(Paths.get(output + "/src/main/java/org/openapitools/client/model/Dog.java"),
+                                     "this.petType = PetTypeEnum.fromValue(\"DogType\");");
+        TestUtils.assertFileContains(Paths.get(output + "/src/main/java/org/openapitools/client/model/Cat.java"),
+                                     "this.petType = PetTypeEnum.fromValue(\"CatType\");");
+
+        output.deleteOnExit();
+    }
 }
diff --git a/modules/openapi-generator/src/test/resources/3_0/issue_806_enum_discriminator.yaml b/modules/openapi-generator/src/test/resources/3_0/issue_806_enum_discriminator.yaml
@@ -0,0 +1,45 @@
+openapi: 3.0.1
+info:
+  title: example api
+  version: v1.0
+paths:
+  /pet:
+    post:
+      requestBody:
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/Pet'
+        required: true
+      responses:
+        "200":
+          description: OK
+components:
+  schemas:
+    Pet:
+      type: object
+      properties:
+        petType:
+          type: string
+          enum:
+          - DogType
+          - CatType
+          - LizardType
+      discriminator:
+        propertyName: petType
+        mapping:
+          DogType: '#/components/schemas/Dog'
+          CatType: '#/components/schemas/Cat'
+          # missing LizardType
+    Dog:
+      type: object
+      allOf:
+      - $ref: '#/components/schemas/Pet'
+    Cat:
+      type: object
+      allOf:
+      - $ref: '#/components/schemas/Pet'
+    Lizard:
+      type: object
+      allOf:
+      - $ref: '#/components/schemas/Pet'
diff --git a/modules/openapi-generator/src/test/resources/3_0/issue_806_string_discriminator.yaml b/modules/openapi-generator/src/test/resources/3_0/issue_806_string_discriminator.yaml
@@ -0,0 +1,41 @@
+openapi: 3.0.1
+info:
+  title: example api
+  version: v1.0
+paths:
+  /pet:
+    post:
+      requestBody:
+        content:
+          application/json:
+            schema:
+              $ref: '#/components/schemas/Pet'
+        required: true
+      responses:
+        "200":
+          description: OK
+components:
+  schemas:
+    Pet:
+      type: object
+      properties:
+        petType:
+          type: string
+      discriminator:
+        propertyName: petType
+        mapping:
+          DogType: '#/components/schemas/Dog'
+          CatType: '#/components/schemas/Cat'
+          # missing LizardType
+    Dog:
+      type: object
+      allOf:
+      - $ref: '#/components/schemas/Pet'
+    Cat:
+      type: object
+      allOf:
+      - $ref: '#/components/schemas/Pet'
+    Lizard:
+      type: object
+      allOf:
+      - $ref: '#/components/schemas/Pet'