Simple oneof support

SuduIDE · Oct 13, 2023 · e6e272d · e6e272d
1 parent 44fabde
commit e6e272d
Show file tree

Hide file tree

Showing 5 changed files with 120 additions and 2 deletions.
diff --git a/generator/src/main/java/org/sudu/protogen/descriptors/Message.java b/generator/src/main/java/org/sudu/protogen/descriptors/Message.java
@@ -85,6 +85,12 @@ public boolean doGenerate() {
         return super.doGenerate();
     }
 
+    public List<OneOf> getOneofs() {
+        return messageDescriptor.getOneofs().stream()
+                .map(o -> new OneOf(o, this))
+                .toList();
+    }
+
     public Optional<String> getTopic() {
         return Options.wrapExtension(messageDescriptor.getOptions(), protogen.Options.topic);
     }

diff --git a/generator/src/main/java/org/sudu/protogen/descriptors/OneOf.java b/generator/src/main/java/org/sudu/protogen/descriptors/OneOf.java
@@ -0,0 +1,58 @@
+package org.sudu.protogen.descriptors;
+
+import com.google.protobuf.Descriptors;
+import com.squareup.javapoet.ClassName;
+import org.jetbrains.annotations.NotNull;
+import org.sudu.protogen.config.naming.NamingManager;
+import org.sudu.protogen.utils.Name;
+
+import java.util.List;
+import java.util.Objects;
+
+public class OneOf implements Descriptor {
+
+    private final Descriptors.OneofDescriptor oneofDescriptor;
+
+    private final Message parent;
+
+    public OneOf(Descriptors.OneofDescriptor oneofDescriptor, Message parent) {
+        this.oneofDescriptor = oneofDescriptor;
+        this.parent = parent;
+    }
+
+    @NotNull
+    public String getName() {
+        return oneofDescriptor.getName();
+    }
+
+    @NotNull
+    public List<String> getFieldsCases() {
+        return oneofDescriptor.getFields().stream()
+                .map(Descriptors.FieldDescriptor::getName)
+                .map(String::toUpperCase)
+                .toList();
+    }
+
+    public ClassName getProtobufTypeName() {
+        ClassName parentClass = parent.getProtobufTypeName();
+        return ClassName.get(parentClass.packageName(), parentClass.simpleName(), Name.toCamelCase(getName()) + "Case");
+    }
+
+    public ClassName getDomainTypeName(NamingManager namingManager) {
+        ClassName parentClass = parent.getDomainTypeName(namingManager);
+        return ClassName.get(parentClass.packageName(), parentClass.simpleName(), Name.toCamelCase(getName()) + "Case");
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        OneOf oneOf = (OneOf) o;
+        return Objects.equals(oneofDescriptor, oneOf.oneofDescriptor) && Objects.equals(parent, oneOf.parent);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(oneofDescriptor, parent);
+    }
+}
diff --git a/generator/src/main/java/org/sudu/protogen/generator/message/MessageGenerator.java b/generator/src/main/java/org/sudu/protogen/generator/message/MessageGenerator.java
@@ -7,6 +7,7 @@
 import org.sudu.protogen.generator.DescriptorGenerator;
 import org.sudu.protogen.generator.GenerationContext;
 import org.sudu.protogen.generator.field.FieldProcessingResult;
+import org.sudu.protogen.utils.Name;
 import org.sudu.protogen.utils.Poem;
 
 import javax.lang.model.element.Modifier;
@@ -33,6 +34,7 @@ public TypeSpec generate(@NotNull Message msgDescriptor) {
         implementComparable(msgDescriptor, typeBuilder);
         addTopicField(msgDescriptor, typeBuilder);
         addTransformingMethods(msgDescriptor, processedFields, typeBuilder);
+        addOneofs(msgDescriptor, typeBuilder);
 
         return typeBuilder
                 .multiLineRecord(true)
@@ -41,6 +43,43 @@ public TypeSpec generate(@NotNull Message msgDescriptor) {
                 .build();
     }
 
+    private void addOneofs(Message msgDescriptor, TypeSpec.Builder typeBuilder) {
+        msgDescriptor.getOneofs().forEach(oneOf -> {
+            if (oneOf.getFieldsCases().size() < 2) return;
+            String oneOfName = Name.toCamelCase(oneOf.getName()) + "Case";
+            ClassName domainTypeName = oneOf.getDomainTypeName(generationContext.configuration().namingManager());
+
+            TypeSpec.Builder oneOfSpecBuilder = TypeSpec.enumBuilder(oneOfName);
+            oneOf.getFieldsCases().forEach(oneOfSpecBuilder::addEnumConstant);
+            oneOfSpecBuilder.addEnumConstant("NOT_SET");
+            oneOfSpecBuilder.addMethod(MethodSpec.methodBuilder("fromProto")
+                    .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
+                    .addParameter(ParameterSpec.builder(oneOf.getProtobufTypeName(), "proto").build())
+                    .addStatement("""
+                            return switch(proto) {$>
+                            $L
+                            case $L -> NOT_SET;
+                            $<}""",
+                            oneOf.getFieldsCases().stream()
+                                    .map(c -> CodeBlock.of("case $L -> $L;", c, c))
+                                    .collect(Poem.joinCodeBlocks("\n")),
+                            oneOf.getName().toUpperCase() + "_NOT_SET"
+                    )
+                    .returns(domainTypeName)
+                    .build()
+            );
+            TypeSpec oneOfSpec = oneOfSpecBuilder.build();
+            typeBuilder.addType(oneOfSpec);
+
+            typeBuilder.addMethod(MethodSpec.methodBuilder("get" + oneOfName)
+                    .addModifiers(Modifier.PUBLIC)
+                    .returns(domainTypeName)
+                    .addStatement("return $T.fromProto(toGrpc().get$L())", domainTypeName, oneOfName)
+                    .build()
+            );
+        });
+    }
+
     private void addComponents(List<FieldProcessingResult> processedFields, TypeSpec.Builder typeBuilder) {
         List<ParameterSpec> parameters = processedFields.stream()
                 .map(FieldProcessingResult::field)

diff --git a/generator/src/main/java/org/sudu/protogen/utils/Poem.java b/generator/src/main/java/org/sudu/protogen/utils/Poem.java
@@ -62,6 +62,10 @@ public static void attachNullabilityAnnotations(
     }
 
     public static Collector<CodeBlock, CodeBlock.Builder, CodeBlock> joinCodeBlocks() {
+        return joinCodeBlocks("");
+    }
+
+    public static Collector<CodeBlock, CodeBlock.Builder, CodeBlock> joinCodeBlocks(String sep) {
         return new Collector<>() {
             @Override
             public Supplier<CodeBlock.Builder> supplier() {
@@ -70,12 +74,15 @@ public Supplier<CodeBlock.Builder> supplier() {
 
             @Override
             public BiConsumer<CodeBlock.Builder, CodeBlock> accumulator() {
-                return CodeBlock.Builder::add;
+                return (builder, block) -> {
+                    if (!builder.isEmpty()) builder.add(sep);
+                    builder.add(block);
+                };
             }
 
             @Override
             public BinaryOperator<CodeBlock.Builder> combiner() {
-                return (builder1, builder2) -> builder1.add(builder2.build());
+                return (builder1, builder2) -> builder1.add(sep).add(builder2.build());
             }
 
             @Override

diff --git a/tests/src/test/proto/general.proto b/tests/src/test/proto/general.proto
@@ -62,4 +62,12 @@ message GrpcRepeatedTypes {
 message GrpcSomeEvent {
   repeated string ids = 1;
   option (.protogen.topic) = "SOME_TOPIC";
+}
+
+message GrpcWithOneofs {
+  string name = 1;
+  oneof named {
+    int32 a = 2;
+    int32 b= 3;
+  }
 }