Transform bytes to byte[] instead of ByteString

SuduIDE · Dec 8, 2023 · 17a712b · 17a712b
1 parent c7d0e9d
commit 17a712b
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 4 deletions.
diff --git a/...c/main/java/org/sudu/protogen/generator/field/processors/PrimitiveFieldTypeProcessor.java b/...c/main/java/org/sudu/protogen/generator/field/processors/PrimitiveFieldTypeProcessor.java
@@ -5,6 +5,7 @@
 import org.jetbrains.annotations.NotNull;
 import org.sudu.protogen.descriptors.Field;
 import org.sudu.protogen.generator.GenerationContext;
+import org.sudu.protogen.generator.type.BytesType;
 import org.sudu.protogen.generator.type.PrimitiveTypeModel;
 import org.sudu.protogen.generator.type.TypeModel;
 
@@ -23,14 +24,12 @@ public PrimitiveFieldTypeProcessor(@NotNull GenerationContext context) {
             case DOUBLE -> boxIfNullable(field, TypeName.DOUBLE);
             case BOOLEAN -> boxIfNullable(field, TypeName.BOOLEAN);
             case STRING -> boxIfNullable(field, ClassName.get(String.class));
-            case BYTE_STRING -> new TypeModel(ClassName.get("com.google.protobuf", "ByteString"));
+            case BYTE_STRING -> new BytesType();
             default -> next(field);
         };
     }
 
     private TypeModel boxIfNullable(@NotNull Field field, TypeName primitive) {
-        return field.isNullable()
-                ? new TypeModel(primitive.box())
-                : new PrimitiveTypeModel(primitive);
+        return field.isNullable() ? new TypeModel(primitive.box()) : new PrimitiveTypeModel(primitive);
     }
 }
diff --git a/generator/src/main/java/org/sudu/protogen/generator/type/BytesType.java b/generator/src/main/java/org/sudu/protogen/generator/type/BytesType.java
@@ -0,0 +1,25 @@
+package org.sudu.protogen.generator.type;
+
+import com.google.protobuf.ByteString;
+import com.squareup.javapoet.ArrayTypeName;
+import com.squareup.javapoet.CodeBlock;
+import com.squareup.javapoet.TypeName;
+
+import java.util.Set;
+
+public class BytesType extends TypeModel {
+
+    public BytesType() {
+        super(ArrayTypeName.of(TypeName.BYTE));
+    }
+
+    @Override
+    public CodeBlock toGrpcTransformer(CodeBlock expr, Set<String> usedDefinitions) {
+        return CodeBlock.of("$T.copyFrom($L)", ByteString.class, expr);
+    }
+
+    @Override
+    public CodeBlock fromGrpcTransformer(CodeBlock expr, Set<String> usedDefinitions) {
+        return CodeBlock.of("$L.toByteArray()", expr);
+    }
+}
diff --git a/tests/src/test/proto/general.proto b/tests/src/test/proto/general.proto
@@ -71,4 +71,8 @@ message GrpcWithOneofs {
     int32 b= 3;
   }
   option (.protogen.builder_for_nullable) = false;
+}
+
+message GrpcByteObject {
+  bytes content = 1;
 }