From 25e1199b28e3b381eb71c416ea1de2581baa2b8c Mon Sep 17 00:00:00 2001 From: Dmitrii Duzhinskii Date: Thu, 26 Oct 2023 13:00:42 +0300 Subject: [PATCH] Add builders for messages with optional fields --- .../sudu/protogen/descriptors/Message.java | 5 + .../protogen/generator/GenerationContext.java | 6 + .../message/MessageBuilderGenerator.java | 103 ++++++++++++++++++ .../generator/message/MessageGenerator.java | 44 +++++++- options/src/main/proto/protogen/options.proto | 1 + tests/src/test/proto/general.proto | 1 + 6 files changed, 156 insertions(+), 4 deletions(-) create mode 100644 generator/src/main/java/org/sudu/protogen/generator/message/MessageBuilderGenerator.java diff --git a/generator/src/main/java/org/sudu/protogen/descriptors/Message.java b/generator/src/main/java/org/sudu/protogen/descriptors/Message.java index 82d56a3..ee572df 100644 --- 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,11 @@ public boolean doGenerate() { return super.doGenerate(); } + public boolean generateBuilderOption() { + return Options.wrapExtension(messageDescriptor.getOptions(), protogen.Options.builderForNullable) + .orElse(true); + } + public List getOneofs() { return messageDescriptor.getOneofs().stream() .map(o -> new OneOf(o, this)) diff --git a/generator/src/main/java/org/sudu/protogen/generator/GenerationContext.java b/generator/src/main/java/org/sudu/protogen/generator/GenerationContext.java index 79b1cf4..ccbbdcc 100644 --- a/generator/src/main/java/org/sudu/protogen/generator/GenerationContext.java +++ b/generator/src/main/java/org/sudu/protogen/generator/GenerationContext.java @@ -9,6 +9,7 @@ import org.sudu.protogen.generator.field.FieldGenerator; import org.sudu.protogen.generator.field.FieldProcessingResult; import org.sudu.protogen.generator.field.processors.*; +import org.sudu.protogen.generator.message.MessageBuilderGenerator; import org.sudu.protogen.generator.message.MessageGenerator; import org.sudu.protogen.generator.server.ServiceGenerator; import org.sudu.protogen.generator.type.TypeModel; @@ -78,6 +79,7 @@ public class GeneratorsHolder { private final DescriptorGenerator fieldGenerator = new FieldGenerator(GenerationContext.this).withCache(); private final DescriptorGenerator messageGenerator = new MessageGenerator(GenerationContext.this).withCache(); + private final DescriptorGenerator messageBuilderGenerator = new MessageBuilderGenerator(GenerationContext.this).withCache(); private final DescriptorGenerator enumGenerator = new EnumGenerator(GenerationContext.this).withCache(); private final DescriptorGenerator clientGenerator = new ClientGenerator(GenerationContext.this).withCache(); private final DescriptorGenerator serviceGenerator = new ServiceGenerator(GenerationContext.this).withCache(); @@ -94,6 +96,10 @@ public TypeSpec generate(Message message) { return messageGenerator.generate(message); } + public TypeSpec generateBuilder(Message message) { + return messageBuilderGenerator.generate(message); + } + public TypeSpec generate(EnumOrMessage enumOrMessage) { if (enumOrMessage instanceof Enum en) return generate(en); if (enumOrMessage instanceof Message msg) return generate(msg); diff --git a/generator/src/main/java/org/sudu/protogen/generator/message/MessageBuilderGenerator.java b/generator/src/main/java/org/sudu/protogen/generator/message/MessageBuilderGenerator.java new file mode 100644 index 0000000..ef45393 --- /dev/null +++ b/generator/src/main/java/org/sudu/protogen/generator/message/MessageBuilderGenerator.java @@ -0,0 +1,103 @@ +package org.sudu.protogen.generator.message; + +import com.squareup.javapoet.*; +import org.jetbrains.annotations.NotNull; +import org.sudu.protogen.descriptors.Message; +import org.sudu.protogen.generator.DescriptorGenerator; +import org.sudu.protogen.generator.GenerationContext; +import org.sudu.protogen.generator.field.FieldGenerationHelper; +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; +import java.util.List; +import java.util.function.Predicate; + +public class MessageBuilderGenerator implements DescriptorGenerator { + + private final GenerationContext context; + + public MessageBuilderGenerator(@NotNull GenerationContext context) { + this.context = context; + } + + @Override + public TypeSpec generate(Message descriptor) { + TypeSpec.Builder builder = TypeSpec.classBuilder(getBuilderName(descriptor)) + .addModifiers(Modifier.PUBLIC, Modifier.STATIC); + addFields(descriptor, builder); + addConstructor(descriptor, builder); + addSetters(descriptor, builder); + addBuildMethod(descriptor, builder); + return builder.build(); + } + + private void addBuildMethod(Message descriptor, TypeSpec.Builder builder) { + ClassName messageType = descriptor.getDomainTypeName(context.configuration().namingManager()); + MethodSpec buildMethod = MethodSpec.methodBuilder("build") + .addModifiers(Modifier.PUBLIC) + .addAnnotation(NotNull.class) + .returns(messageType) + .addStatement("return new $T($L)", messageType, + Poem.separatedSequence( + FieldGenerationHelper.processAllFields(descriptor, context) + .map(FieldProcessingResult::field) + .map(field -> CodeBlock.of("$N", field)) + .toList(), + ", " + )) + .build(); + builder.addMethod(buildMethod); + } + + private void addSetters(Message descriptor, TypeSpec.Builder builder) { + String builderName = getBuilderName(descriptor); + FieldGenerationHelper.processAllFields(descriptor, context) + .filter(FieldProcessingResult::isNullable) + .map(FieldProcessingResult::field) + .map(field -> setterForField(field, builderName)) + .forEach(builder::addMethod); + } + + @NotNull + private MethodSpec setterForField(FieldSpec fieldSpec, String builderName) { + return MethodSpec.methodBuilder("set" + Name.toCamelCase(fieldSpec.name)) + .addModifiers(Modifier.PUBLIC) + .returns(ClassName.get("",builderName)) + .addStatement("this.$N = $N", fieldSpec, fieldSpec) + .addStatement("return this") + .build(); + } + + private void addConstructor(Message descriptor, TypeSpec.Builder builder) { + List nonNullFields = FieldGenerationHelper.processAllFields(descriptor, context) + .filter(Predicate.not(FieldProcessingResult::isNullable)) + .map(FieldProcessingResult::field) + .map(field -> field.toBuilder().clearAnnotations().build()) + .toList(); + MethodSpec constructor = MethodSpec.constructorBuilder().addParameters( + nonNullFields.stream() + .map(Poem::fieldToParameter) + .toList() + ).addModifiers(Modifier.PRIVATE).addCode( + nonNullFields.stream() + .map(field -> CodeBlock.of("this.$N = $N;", field, field)) + .collect(Poem.joinCodeBlocks("\n")) + ).build(); + builder.addMethod(constructor); + } + + private void addFields(Message descriptor, TypeSpec.Builder builder) { + FieldGenerationHelper.processAllFields(descriptor, context) + .map(FieldProcessingResult::field) + .map(field -> field.toBuilder().clearAnnotations().addModifiers(Modifier.PRIVATE).build()) + .forEach(builder::addField); + } + + @NotNull + private String getBuilderName(Message descriptor) { + return descriptor.getDomainTypeName(context.configuration().namingManager()).simpleName() + "Builder"; + } + +} 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 index c6b32e2..ad8ebb4 100644 --- a/generator/src/main/java/org/sudu/protogen/generator/message/MessageGenerator.java +++ b/generator/src/main/java/org/sudu/protogen/generator/message/MessageGenerator.java @@ -6,12 +6,15 @@ import org.sudu.protogen.descriptors.Message; import org.sudu.protogen.generator.DescriptorGenerator; import org.sudu.protogen.generator.GenerationContext; +import org.sudu.protogen.generator.field.FieldGenerationHelper; 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; import java.util.List; +import java.util.function.Predicate; +import java.util.stream.Collectors; public class MessageGenerator implements DescriptorGenerator { @@ -35,6 +38,7 @@ public TypeSpec generate(@NotNull Message msgDescriptor) { addTopicField(msgDescriptor, typeBuilder); addTransformingMethods(msgDescriptor, processedFields, typeBuilder); addOneofs(msgDescriptor, typeBuilder); + addBuilderIfNecessary(msgDescriptor, typeBuilder); return typeBuilder .multiLineRecord(true) @@ -43,6 +47,38 @@ public TypeSpec generate(@NotNull Message msgDescriptor) { .build(); } + private void addBuilderIfNecessary(Message msgDescriptor, TypeSpec.Builder typeBuilder) { + if (!msgDescriptor.generateBuilderOption()) return; + List processedFields = FieldGenerationHelper.processAllFields(msgDescriptor, generationContext).toList(); + if (processedFields.stream().anyMatch(FieldProcessingResult::isNullable)) { + typeBuilder.addType(generationContext.generatorsHolder().generateBuilder(msgDescriptor)); + List notNullFields = processedFields.stream().filter(Predicate.not(FieldProcessingResult::isNullable)).map(FieldProcessingResult::field).toList(); + ClassName domainTypeName = msgDescriptor.getDomainTypeName(generationContext.configuration().namingManager()); + ClassName builderType = ClassName.get(domainTypeName.canonicalName(), domainTypeName.simpleName() + "Builder"); + typeBuilder.addMethod(MethodSpec.methodBuilder("builder") + .addModifiers(Modifier.PUBLIC, Modifier.STATIC) + .addAnnotation(NotNull.class) + .returns(builderType) + .addParameters(notNullFields.stream().map(Poem::fieldToParameter).toList()) + .addStatement("return new $T($L)", + builderType, + Poem.separatedSequence(notNullFields.stream().map(f -> CodeBlock.of("$N", f)).toList(), ",") + ) + .build() + ); + String constructorParams = processedFields.stream() + .map(f -> f.isNullable() ? "null" : f.field().name) + .collect(Collectors.joining(", ")); + typeBuilder.addMethod( + MethodSpec.constructorBuilder() + .addModifiers(Modifier.PUBLIC) + .addParameters(notNullFields.stream().map(Poem::fieldToParameter).toList()) + .addStatement("this($L)", constructorParams) + .build() + ); + } + } + private void addOneofs(Message msgDescriptor, TypeSpec.Builder typeBuilder) { msgDescriptor.getOneofs().forEach(oneOf -> { if (oneOf.getFieldsCases().size() < 2) return; @@ -56,10 +92,10 @@ private void addOneofs(Message msgDescriptor, TypeSpec.Builder typeBuilder) { .addModifiers(Modifier.PUBLIC, Modifier.STATIC) .addParameter(ParameterSpec.builder(oneOf.getProtobufTypeName(), "proto").build()) .addStatement(""" - return switch(proto) {$> - $L - case $L -> NOT_SET; - $<}""", + return switch(proto) {$> + $L + case $L -> NOT_SET; + $<}""", oneOf.getFieldsCases().stream() .map(c -> CodeBlock.of("case $L -> $L;", c, c)) .collect(Poem.joinCodeBlocks("\n")), diff --git a/options/src/main/proto/protogen/options.proto b/options/src/main/proto/protogen/options.proto index 318dbbd..b811e7d 100644 --- a/options/src/main/proto/protogen/options.proto +++ b/options/src/main/proto/protogen/options.proto @@ -77,6 +77,7 @@ extend google.protobuf.MessageOptions { */ string message_comparator = 5105; string topic = 5016; + bool builder_for_nullable = 5017; } extend google.protobuf.EnumOptions { diff --git a/tests/src/test/proto/general.proto b/tests/src/test/proto/general.proto index babf5f5..394d8f2 100644 --- a/tests/src/test/proto/general.proto +++ b/tests/src/test/proto/general.proto @@ -70,4 +70,5 @@ message GrpcWithOneofs { int32 a = 2; int32 b= 3; } + option (.protogen.builder_for_nullable) = false; } \ No newline at end of file