cmd/protoc-gen-go: generate enums
This produces exactly the same output (to the best of my ability to
determine) as github.com/golang/protobuf.
Change-Id: Ib60e7a836efb1eb0e5167b30458049ec239e7903
Reviewed-on: https://go-review.googlesource.com/134695
Reviewed-by: Joe Tsai <thebrokentoaster@gmail.com>
diff --git a/cmd/protoc-gen-go/main.go b/cmd/protoc-gen-go/main.go
index d330941..d91e7a1 100644
--- a/cmd/protoc-gen-go/main.go
+++ b/cmd/protoc-gen-go/main.go
@@ -18,8 +18,11 @@
"github.com/golang/protobuf/proto"
descpb "github.com/golang/protobuf/protoc-gen-go/descriptor"
"google.golang.org/proto/protogen"
+ "google.golang.org/proto/reflect/protoreflect"
)
+const protoPackage = "github.com/golang/protobuf/proto"
+
func main() {
protogen.Run(func(gen *protogen.Plugin) error {
for _, f := range gen.Files {
@@ -34,7 +37,9 @@
type File struct {
*protogen.File
- locationMap map[string][]*descpb.SourceCodeInfo_Location
+ locationMap map[string][]*descpb.SourceCodeInfo_Location
+ descriptorVar string // var containing the gzipped FileDescriptorProto
+ init []string
}
func genFile(gen *protogen.Plugin, file *protogen.File) {
@@ -47,6 +52,12 @@
f.locationMap[key] = append(f.locationMap[key], loc)
}
+ // Determine the name of the var holding the file descriptor:
+ //
+ // fileDescriptor_<hash of filename>
+ filenameHash := sha256.Sum256([]byte(f.Desc.Path()))
+ f.descriptorVar = fmt.Sprintf("fileDescriptor_%s", hex.EncodeToString(filenameHash[:8]))
+
g := gen.NewGeneratedFile(f.GeneratedFilenamePrefix+".pb.go", f.GoImportPath)
g.P("// Code generated by protoc-gen-go. DO NOT EDIT.")
g.P("// source: ", f.Desc.Path())
@@ -57,20 +68,26 @@
g.P("package ", f.GoPackageName)
g.P()
+ for _, enum := range f.Enums {
+ genEnum(gen, g, f, enum)
+ }
for _, message := range f.Messages {
genMessage(gen, g, f, message)
}
+ if len(f.init) != 0 {
+ g.P("func init() {")
+ for _, s := range f.init {
+ g.P(s)
+ }
+ g.P("}")
+ g.P()
+ }
+
genFileDescriptor(gen, g, f)
}
func genFileDescriptor(gen *protogen.Plugin, g *protogen.GeneratedFile, f *File) {
- // Determine the name of the var holding the file descriptor:
- //
- // fileDescriptor_<hash of filename>
- filenameHash := sha256.Sum256([]byte(f.Desc.Path()))
- varName := fmt.Sprintf("fileDescriptor_%s", hex.EncodeToString(filenameHash[:8]))
-
// Trim the source_code_info from the descriptor.
// Marshal and gzip it.
descProto := proto.Clone(f.Proto).(*descpb.FileDescriptorProto)
@@ -86,9 +103,9 @@
w.Close()
b = buf.Bytes()
- g.P("func init() { proto.RegisterFile(", strconv.Quote(f.Desc.Path()), ", ", varName, ") }")
+ g.P("func init() { proto.RegisterFile(", strconv.Quote(f.Desc.Path()), ", ", f.descriptorVar, ") }")
g.P()
- g.P("var ", varName, " = []byte{")
+ g.P("var ", f.descriptorVar, " = []byte{")
g.P("// ", len(b), " bytes of a gzipped FileDescriptorProto")
for len(b) > 0 {
n := 16
@@ -108,7 +125,92 @@
g.P()
}
+func genEnum(gen *protogen.Plugin, g *protogen.GeneratedFile, f *File, enum *protogen.Enum) {
+ genComment(g, f, enum.Path)
+ // TODO: deprecation
+ g.P("type ", enum.GoIdent, " int32")
+ g.P("const (")
+ for _, value := range enum.Values {
+ genComment(g, f, value.Path)
+ // TODO: deprecation
+ g.P(value.GoIdent, " ", enum.GoIdent, " = ", value.Desc.Number())
+ }
+ g.P(")")
+ g.P()
+ nameMap := enum.GoIdent.GoName + "_name"
+ g.P("var ", nameMap, " = map[int32]string{")
+ generated := make(map[protoreflect.EnumNumber]bool)
+ for _, value := range enum.Values {
+ duplicate := ""
+ if _, present := generated[value.Desc.Number()]; present {
+ duplicate = "// Duplicate value: "
+ }
+ g.P(duplicate, value.Desc.Number(), ": ", strconv.Quote(string(value.Desc.Name())), ",")
+ generated[value.Desc.Number()] = true
+ }
+ g.P("}")
+ g.P()
+ valueMap := enum.GoIdent.GoName + "_value"
+ g.P("var ", valueMap, " = map[string]int32{")
+ for _, value := range enum.Values {
+ g.P(strconv.Quote(string(value.Desc.Name())), ": ", value.Desc.Number(), ",")
+ }
+ g.P("}")
+ g.P()
+ if enum.Desc.Syntax() != protoreflect.Proto3 {
+ g.P("func (x ", enum.GoIdent, ") Enum() *", enum.GoIdent, " {")
+ g.P("p := new(", enum.GoIdent, ")")
+ g.P("*p = x")
+ g.P("return p")
+ g.P("}")
+ g.P()
+ }
+ g.P("func (x ", enum.GoIdent, ") String() string {")
+ g.P("return ", protogen.GoIdent{GoImportPath: protoPackage, GoName: "EnumName"}, "(", enum.GoIdent, "_name, int32(x))")
+ g.P("}")
+ g.P()
+
+ if enum.Desc.Syntax() != protoreflect.Proto3 {
+ g.P("func (x *", enum.GoIdent, ") UnmarshalJSON(data []byte) error {")
+ g.P("value, err := ", protogen.GoIdent{GoImportPath: protoPackage, GoName: "UnmarshalJSONEnum"}, "(", enum.GoIdent, `_value, data, "`, enum.GoIdent, `")`)
+ g.P("if err != nil {")
+ g.P("return err")
+ g.P("}")
+ g.P("*x = ", enum.GoIdent, "(value)")
+ g.P("return nil")
+ g.P("}")
+ g.P()
+ }
+
+ var indexes []string
+ for i := 1; i < len(enum.Path); i += 2 {
+ indexes = append(indexes, strconv.Itoa(int(enum.Path[i])))
+ }
+ g.P("func (", enum.GoIdent, ") EnumDescriptor() ([]byte, []int) {")
+ g.P("return ", f.descriptorVar, ", []int{", strings.Join(indexes, ","), "}")
+ g.P("}")
+ g.P()
+
+ genWellKnownType(g, enum.GoIdent, enum.Desc)
+
+ // The name registered is, confusingly, <proto_package>.<go_ident>.
+ // This probably should have been the full name of the proto enum
+ // type instead, but changing it at this point would require thought.
+ regName := string(f.Desc.Package()) + "." + enum.GoIdent.GoName
+ f.init = append(f.init, fmt.Sprintf("%s(%q, %s, %s)",
+ g.QualifiedGoIdent(protogen.GoIdent{
+ GoImportPath: protoPackage,
+ GoName: "RegisterEnum",
+ }),
+ regName, nameMap, valueMap,
+ ))
+}
+
func genMessage(gen *protogen.Plugin, g *protogen.GeneratedFile, f *File, message *protogen.Message) {
+ for _, enum := range message.Enums {
+ genEnum(gen, g, f, enum)
+ }
+
genComment(g, f, message.Path)
g.P("type ", message.GoIdent, " struct {")
g.P("}")
@@ -142,3 +244,15 @@
}
return string(buf)
}
+
+func genWellKnownType(g *protogen.GeneratedFile, ident protogen.GoIdent, desc protoreflect.Descriptor) {
+ if wellKnownTypes[desc.FullName()] {
+ g.P("func (", ident, `) XXX_WellKnownType() string { return "`, desc.Name(), `" }`)
+ g.P()
+ }
+}
+
+// Names of messages and enums for which we will generate XXX_WellKnownType methods.
+var wellKnownTypes = map[protoreflect.FullName]bool{
+ "google.protobuf.NullValue": true,
+}