Skip to content

Commit

Permalink
Add builders for messages with optional fields
Browse files Browse the repository at this point in the history
  • Loading branch information
Duzhinsky committed Oct 26, 2023
1 parent 200759e commit 25e1199
Show file tree
Hide file tree
Showing 6 changed files with 156 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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<OneOf> getOneofs() {
return messageDescriptor.getOneofs().stream()
.map(o -> new OneOf(o, this))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -78,6 +79,7 @@ public class GeneratorsHolder {

private final DescriptorGenerator<Field, FieldProcessingResult> fieldGenerator = new FieldGenerator(GenerationContext.this).withCache();
private final DescriptorGenerator<Message, TypeSpec> messageGenerator = new MessageGenerator(GenerationContext.this).withCache();
private final DescriptorGenerator<Message, TypeSpec> messageBuilderGenerator = new MessageBuilderGenerator(GenerationContext.this).withCache();
private final DescriptorGenerator<Enum, TypeSpec> enumGenerator = new EnumGenerator(GenerationContext.this).withCache();
private final DescriptorGenerator<Service, TypeSpec> clientGenerator = new ClientGenerator(GenerationContext.this).withCache();
private final DescriptorGenerator<Service, TypeSpec> serviceGenerator = new ServiceGenerator(GenerationContext.this).withCache();
Expand All @@ -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);
Expand Down
Original file line number Diff line number Diff line change
@@ -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<Message, TypeSpec> {

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<FieldSpec> 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";
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -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<Message, TypeSpec> {

Expand All @@ -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)
Expand All @@ -43,6 +47,38 @@ public TypeSpec generate(@NotNull Message msgDescriptor) {
.build();
}

private void addBuilderIfNecessary(Message msgDescriptor, TypeSpec.Builder typeBuilder) {
if (!msgDescriptor.generateBuilderOption()) return;
List<FieldProcessingResult> processedFields = FieldGenerationHelper.processAllFields(msgDescriptor, generationContext).toList();
if (processedFields.stream().anyMatch(FieldProcessingResult::isNullable)) {
typeBuilder.addType(generationContext.generatorsHolder().generateBuilder(msgDescriptor));
List<FieldSpec> 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;
Expand All @@ -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")),
Expand Down
1 change: 1 addition & 0 deletions options/src/main/proto/protogen/options.proto
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,7 @@ extend google.protobuf.MessageOptions {
*/
string message_comparator = 5105;
string topic = 5016;
bool builder_for_nullable = 5017;
}

extend google.protobuf.EnumOptions {
Expand Down
1 change: 1 addition & 0 deletions tests/src/test/proto/general.proto
Original file line number Diff line number Diff line change
Expand Up @@ -70,4 +70,5 @@ message GrpcWithOneofs {
int32 a = 2;
int32 b= 3;
}
option (.protogen.builder_for_nullable) = false;
}

0 comments on commit 25e1199

Please sign in to comment.