small refactor in ProtoProperty

eclipse-vertx · Sep 13, 2023 · 4dbeab8 · 4dbeab8
1 parent 8d5e333
commit 4dbeab8
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 13 deletions.
diff --git a/vertx-codegen-protobuf/src/main/java/io/vertx/codegen/protobuf/generator/ProtoFileGen.java b/vertx-codegen-protobuf/src/main/java/io/vertx/codegen/protobuf/generator/ProtoFileGen.java
@@ -97,23 +97,27 @@ private String renderDataObjectModel(DataObjectModel model, int index) {
       ClassKind propKind = prop.getType().getKind();
       ProtoProperty protoProperty = ProtoProperty.getProtoProperty(prop, fieldNumber);
 
-      String protoType;
+      String protoFieldType;
       if (propKind.basic) {
-        protoType = protoProperty.getProtoType().value;
+        protoFieldType = protoProperty.getProtoType().value;
       } else {
-        if (protoProperty.isBuiltinType()) {
-          protoType = "io.vertx.protobuf." + protoProperty.getBuiltInType();
-        } else {
-          protoType = protoProperty.getMessage();
+        if (prop.getType().getKind() == ClassKind.ENUM) {
+            protoFieldType = protoProperty.getEnumType();
+        } else { // Not Enum
+          if (protoProperty.isBuiltinType()) {
+            protoFieldType = "io.vertx.protobuf." + protoProperty.getBuiltInType();
+          } else {
+            protoFieldType = protoProperty.getMessage();
+          }
         }
       }
 
       if (prop.getKind().isList()) {
-        writer.print("  repeated " + protoType + " " + prop.getName() + " = " + fieldNumber + ";\n");
+        writer.print("  repeated " + protoFieldType + " " + prop.getName() + " = " + fieldNumber + ";\n");
       } else if (prop.getKind().isMap()) {
-        writer.print("  map<string, " + protoType + "> " + prop.getName() + " = " + fieldNumber + ";\n");
+        writer.print("  map<string, " + protoFieldType + "> " + prop.getName() + " = " + fieldNumber + ";\n");
       } else {
-        writer.print("  " + protoType + " " + prop.getName() + " = " + fieldNumber + ";\n");
+        writer.print("  " + protoFieldType + " " + prop.getName() + " = " + fieldNumber + ";\n");
       }
       fieldNumber++;
     }

diff --git a/vertx-codegen-protobuf/src/main/java/io/vertx/codegen/protobuf/generator/ProtoProperty.java b/vertx-codegen-protobuf/src/main/java/io/vertx/codegen/protobuf/generator/ProtoProperty.java
@@ -15,21 +15,23 @@ public class ProtoProperty {
   private boolean isNullable;
   // Indicate the name of the message of the nested field
   private String message;
+  // Indicate the name of the enumeration type
+  private String enumType;
   // Built-in types are predefined complex proto types
   // Examples: datetime.proto, struct.proto, vertx-struct.proto
   private String builtInType;
 
   public static ProtoProperty getProtoProperty(PropertyInfo prop, int fieldNumber) {
     ProtoProperty protoProperty = new ProtoProperty();
     ClassKind propKind = prop.getType().getKind();
-    ProtoType protoType;
+    ProtoType protoType = null;
     boolean isNullable = determineIsNullable(prop.getType().getName());
     String message = null;
+    String enumType = null;
     String builtInProtoType = null;
     int wireType;
     if (prop.getType().getKind() == ClassKind.ENUM) {
-      protoType = null;
-      message = prop.getType().getSimpleName(); // TODO Should not call 'message'
+      enumType = prop.getType().getSimpleName();
       wireType = 0;
     } else { // Not Enum
       if (propKind.basic) {
@@ -53,7 +55,6 @@ public static ProtoProperty getProtoProperty(PropertyInfo prop, int fieldNumber)
             throw new UnsupportedOperationException("Unsupported proto-type " + protoType);
         }
       } else {
-        protoType = null;
         message = prop.getType().getSimpleName();
         wireType = 2;
         builtInProtoType = determineBuiltInType(prop);
@@ -74,6 +75,7 @@ public static ProtoProperty getProtoProperty(PropertyInfo prop, int fieldNumber)
     protoProperty.tag = tag;
     protoProperty.protoType = protoType;
     protoProperty.isNullable = isNullable;
+    protoProperty.enumType = enumType;
     protoProperty.message = message;
     protoProperty.builtInType = builtInProtoType;
     return protoProperty;
@@ -175,6 +177,10 @@ public boolean isNullable() {
     return isNullable;
   }
 
+  public String getEnumType() {
+    return enumType;
+  }
+
   public String getMessage() {
     return message;
   }