add extension generations

Author: joesonw

example/output.proto

@@ -5,23 +5,32 @@ option go_package = 'generated/proto;pb';
 
 package example.com;
 
+import 'google/protobuf/descriptor.proto';
+
+extend google.protobuf.MessageOptions {
+    string test = 3001;
+    string abc = 3002;
+}
+
 message Base {
     option (test) = 'hello';
 
     int64 id = 1;
+
 }
 
-// extends example.com.Base
+
 message Intermediate {
     option (test) = 'hello';
 
     int64 id = 1;
     int64 createdAt = 2;
     int64 updatedAt = 3;
+
 }
 
+
 message groups {
-    // extends example.com.Intermediate
     message Group {
         option (test) = 'hello';
 
@@ -30,12 +39,13 @@ message groups {
         int64 updatedAt = 3;
         string name = 11;
         repeated users.User users = 12;
+
     }
 
+
 }
 
 message users {
-    // extends example.com.Intermediate
     message User {
         option (test) = 'hello';
 
@@ -46,19 +56,22 @@ message users {
         string email = 12;
         string password = 13;
         repeated groups.Group groups = 14;
+
     }
 
+
     message nested {
-        // extends example.com.Intermediate
         message Test {
             option (test) = 'hello';
 
             int64 id = 1;
             int64 createdAt = 2;
             int64 updatedAt = 3;
             string name = 11;
+
         }
 
+
     }
 
 }

example/proto/annotation.proto

@@ -7,4 +7,8 @@ import 'google/protobuf/descriptor.proto';
 
 extend google.protobuf.MessageOptions {
     string test = 3001;
+}
+
+extend google.protobuf.MessageOptions {
+    string abc = 3002;
 }

example/template.proto

@@ -4,3 +4,5 @@ syntax = 'proto3';
 option go_package = 'generated/proto;pb';
 
 package example.com;
+
+import 'google/protobuf/descriptor.proto';

main.go

@@ -14,7 +14,7 @@ import (
 	"github.com/spf13/pflag"
 	"google.golang.org/protobuf/types/descriptorpb"
 
-	oneprotou_til "github.com/joesonw/oneproto/util"
+	oneproto_util "github.com/joesonw/oneproto/util"
 )
 
 var (
@@ -71,7 +71,7 @@ func main() {
 		log.Fatalln(err)
 	}
 
-	buf := oneprotou_til.NewBuffer()
+	buf := oneproto_util.NewBuffer()
 	buf.Printf(string(templateContent))
 
 	// group files by package name
@@ -115,20 +115,23 @@ func main() {
 
 		for _, file := range group.files {
 			for _, enum := range file.EnumType {
-				oneprotou_til.GenerateEnum(buf, identLevel+1, enum)
+				oneproto_util.GenerateEnum(buf, identLevel+1, enum)
 				buf.Printf("")
 			}
 
 			for _, service := range file.Service {
-				oneprotou_til.GenerateService(buf, identLevel+1, service)
+				oneproto_util.GenerateService(buf, identLevel+1, service)
 				buf.Printf("")
 			}
 
 			for _, message := range file.GetMessageType() {
-				resolveMessageExtends(buf, identLevel+1, message)
-				oneprotou_til.GenerateMessage(buf, identLevel+1, message)
+				resolveMessageExtends(message)
+				oneproto_util.GenerateMessage(buf, identLevel+1, message)
 				buf.Printf("")
 			}
+
+			oneproto_util.GenerateExtensions(buf, identLevel+1, nil, file.GetExtension())
+			buf.Printf("")
 		}
 
 		for _, subPackage := range group.subPackages {
@@ -146,29 +149,42 @@ func main() {
 	}
 }
 
-func resolveMessageExtends(buf *oneprotou_til.Buffer, indentLevel int, message *descriptorpb.DescriptorProto) {
+func resolveMessageExtends(message *descriptorpb.DescriptorProto) {
 	if parentResolvedMap[message] {
 		return
 	}
+	for _, nested := range message.GetNestedType() {
+		resolveMessageExtends(nested)
+	}
+
 	parentResolvedMap[message] = true
+	var options []*descriptorpb.UninterpretedOption
+	var fields []*descriptorpb.FieldDescriptorProto
 	for i, option := range message.GetOptions().GetUninterpretedOption() {
 		if isOptionOneProtoExtends(option) {
 			message.Options.UninterpretedOption = append(message.Options.UninterpretedOption[:i], message.Options.UninterpretedOption[i+1:]...)
-			buf.Printf("%s// extends %s", strings.Repeat(" ", 4*indentLevel), option.GetStringValue())
 			parent := allMessageDescriptors[trimPackageFromName(string(option.GetStringValue()))]
 			if parent == nil {
 				log.Fatalf("unable to find message %s", option.GetStringValue())
 			}
-			resolveMessageExtends(buf, indentLevel, parent)
-			message.Field = append(message.Field, parent.Field...)
-			if message.Options == nil {
-				message.Options = &descriptorpb.MessageOptions{}
-			}
-			if *pOptions {
-				message.Options.UninterpretedOption = append(message.Options.UninterpretedOption, parent.GetOptions().GetUninterpretedOption()...)
+			resolveMessageExtends(parent)
+			fields = append(fields, parent.Field...)
+			options = append(options, parent.GetOptions().GetUninterpretedOption()...)
+		}
+	}
+
+	message.Field = append(message.Field, fields...)
+	if *pOptions {
+		if message.Options == nil {
+			message.Options = &descriptorpb.MessageOptions{}
+		}
+		for i := range options {
+			if isOptionOneProtoExtends(options[i]) {
+				options = append(options[:i], options[i+1:]...)
 			}
 		}
 	}
+	message.Options.UninterpretedOption = append(message.Options.UninterpretedOption, options...)
 	sort.Slice(message.Field, func(i, j int) bool {
 		return message.Field[i].GetNumber() < message.Field[j].GetNumber()
 	})

util/buffer.go

@@ -24,3 +24,15 @@ func (b *Buffer) String() string {
 func (b *Buffer) Bytes() []byte {
 	return b.buf.Bytes()
 }
+
+func (b *Buffer) Close() error {
+	return nil
+}
+
+func (b *Buffer) Write(p []byte) (int, error) {
+	return b.buf.Write(p)
+}
+
+func (b *Buffer) Read(p []byte) (int, error) {
+	return b.buf.Read(p)
+}

util/generate.go

@@ -9,7 +9,7 @@ import (
 func GenerateService(buf *Buffer, indentLevel int, service *descriptorpb.ServiceDescriptorProto) {
 	indent := strings.Repeat(" ", indentLevel*4)
 	buf.Printf("%sservice %s {", indent, service.GetName())
-	GenerateHeadOptions(buf, indent, service.GetOptions().GetUninterpretedOption())
+	generateHeadOptions(buf, indent, service.GetOptions().GetUninterpretedOption())
 	for _, method := range service.Method {
 		buf.Printf("%s    rpc %s(%s) returns (%s) {", indent, method.GetName(), method.GetInputType(), method.GetOutputType())
 		for _, opt := range method.GetOptions().GetUninterpretedOption() {
@@ -24,7 +24,7 @@ func GenerateService(buf *Buffer, indentLevel int, service *descriptorpb.Service
 func GenerateEnum(buf *Buffer, indentLevel int, enum *descriptorpb.EnumDescriptorProto) {
 	indent := strings.Repeat(" ", indentLevel*4)
 	buf.Printf("%senum %s {", indent, enum.GetName())
-	GenerateHeadOptions(buf, indent, enum.GetOptions().GetUninterpretedOption())
+	generateHeadOptions(buf, indent, enum.GetOptions().GetUninterpretedOption())
 
 	for _, value := range enum.Value {
 		buf.Printf("%s    %s  = %d%s;", indent, value.GetName(), value.GetNumber(), StringifyValueOptions(value.GetOptions().GetUninterpretedOption()))
@@ -35,27 +35,30 @@ func GenerateEnum(buf *Buffer, indentLevel int, enum *descriptorpb.EnumDescripto
 func GenerateMessage(buf *Buffer, indentLevel int, message *descriptorpb.DescriptorProto) {
 	indent := strings.Repeat(" ", indentLevel*4)
 	buf.Printf("%smessage %s {", indent, message.GetName())
-	GenerateHeadOptions(buf, indent, message.GetOptions().GetUninterpretedOption())
+	generateHeadOptions(buf, indent, message.GetOptions().GetUninterpretedOption())
 	for _, field := range message.GetField() {
 		buf.Printf("%s    %s %s = %d%s;", indent, StringifyField(message, field), field.GetName(), field.GetNumber(), StringifyValueOptions(field.GetOptions().GetUninterpretedOption()))
 	}
 
-	for _, enum := range message.EnumType {
+	for _, enum := range message.GetEnumType() {
 		buf.Printf("")
 		GenerateEnum(buf, indentLevel+1, enum)
 	}
 
-	for _, nested := range message.NestedType {
+	for _, nested := range message.GetNestedType() {
 		if nested.GetOptions().GetMapEntry() {
 			continue
 		}
 		buf.Printf("")
 		GenerateMessage(buf, indentLevel+1, nested)
 	}
+
+	buf.Printf("")
+	GenerateExtensions(buf, indentLevel+1, message, message.GetExtension())
 	buf.Printf("%s}", indent)
 }
 
-func GenerateHeadOptions(buf *Buffer, indent string, options []*descriptorpb.UninterpretedOption) {
+func generateHeadOptions(buf *Buffer, indent string, options []*descriptorpb.UninterpretedOption) {
 	if len(options) == 0 {
 		return
 	}
@@ -64,3 +67,19 @@ func GenerateHeadOptions(buf *Buffer, indent string, options []*descriptorpb.Uni
 	}
 	buf.Printf("")
 }
+
+func GenerateExtensions(buf *Buffer, indentLevel int, message *descriptorpb.DescriptorProto, extensions []*descriptorpb.FieldDescriptorProto) {
+	grouped := map[string][]*descriptorpb.FieldDescriptorProto{}
+	for _, ext := range extensions {
+		grouped[ext.GetExtendee()] = append(grouped[ext.GetExtendee()], ext)
+	}
+	indent := strings.Repeat(" ", indentLevel*4)
+	for extendee, exts := range grouped {
+
+		buf.Printf("%sextend %s {", indent, extendee)
+		for _, ext := range exts {
+			buf.Printf("%s    %s %s = %d%s;", indent, StringifyField(message, ext), ext.GetName(), ext.GetNumber(), StringifyValueOptions(ext.GetOptions().GetUninterpretedOption()))
+		}
+		buf.Printf("%s}", indent)
+	}
+}

util/stringify.go

@@ -33,7 +33,7 @@ func StringifyField(message *descriptorpb.DescriptorProto, field *descriptorpb.F
 	repeated := field.GetLabel() == descriptorpb.FieldDescriptorProto_LABEL_REPEATED
 	if field.Type == nil {
 		name := field.GetTypeName()
-		for _, nested := range message.NestedType {
+		for _, nested := range message.GetNestedType() {
 			if name == nested.GetName() && strings.HasSuffix(name, "Entry") { // map entry
 				return fmt.Sprintf("map<%s,%s>", StringifyField(nested, nested.Field[0]), StringifyField(nested, nested.Field[1]))
 			}