diff --git a/ast.go b/ast.go new file mode 100644 index 0000000..b6adee3 --- /dev/null +++ b/ast.go @@ -0,0 +1,153 @@ +package main + +import ( + "go/ast" + "go/format" + "go/parser" + "go/token" + "os" + "path/filepath" + "strings" + + "github.com/fatih/structtag" + tag_options "github.com/unistack-org/micro-proto/tag" + "google.golang.org/protobuf/compiler/protogen" + "google.golang.org/protobuf/proto" +) + +var ( + astFields = make(map[string]map[string]map[string]*structtag.Tags) // map proto file with proto message ast struct +) + +func (g *Generator) astGenerate(plugin *protogen.Plugin) error { + for _, file := range plugin.Files { + if !file.Generate { + continue + } + + for _, message := range file.Messages { + for _, field := range message.Fields { + if field.Desc.Options() == nil { + continue + } + if !proto.HasExtension(field.Desc.Options(), tag_options.E_Tags) { + continue + } + + opts := proto.GetExtension(field.Desc.Options(), tag_options.E_Tags) + if opts != nil { + fpath := filepath.Join(g.tagPath, file.GeneratedFilenamePrefix+".pb.go") + mp, ok := astFields[fpath] + if !ok { + mp = make(map[string]map[string]*structtag.Tags) + } + nmp := make(map[string]*structtag.Tags) + tags, err := structtag.Parse(opts.(string)) + if err != nil { + return err + } + nmp[field.GoName] = tags + mp[message.GoIdent.GoName] = nmp + astFields[fpath] = mp + } + } + } + } + + for file, mp := range astFields { + fset := token.NewFileSet() + pf, err := parser.ParseFile(fset, file, nil, parser.AllErrors|parser.ParseComments) + if err != nil { + return err + } + + r := retag{} + f := func(n ast.Node) ast.Visitor { + if r.err != nil { + return nil + } + + if v, ok := n.(*ast.TypeSpec); ok { + r.fields = mp[v.Name.Name] + return r + } + + return nil + } + + ast.Walk(structVisitor{f}, pf) + + if r.err != nil { + return err + } + + fp, err := os.OpenFile(file, os.O_WRONLY|os.O_TRUNC, os.FileMode(0644)) + if err != nil { + return err + } + if err = format.Node(fp, fset, pf); err != nil { + fp.Close() + return err + } + if err = fp.Close(); err != nil { + return err + } + } + + return nil +} + +type retag struct { + err error + fields map[string]*structtag.Tags +} + +func (v retag) Visit(n ast.Node) ast.Visitor { + if v.err != nil { + return nil + } + if f, ok := n.(*ast.Field); ok { + if len(f.Names) == 0 { + return nil + } + + newTags := v.fields[f.Names[0].String()] + if newTags == nil { + return nil + } + if f.Tag == nil { + f.Tag = &ast.BasicLit{ + Kind: token.STRING, + } + } + + oldTags, err := structtag.Parse(strings.Trim(f.Tag.Value, "`")) + if err != nil { + v.err = err + return nil + } + for _, t := range newTags.Tags() { + oldTags.Set(t) + } + + f.Tag.Value = "`" + oldTags.String() + "`" + + return nil + } + + return v +} + +type structVisitor struct { + visitor func(n ast.Node) ast.Visitor +} + +func (v structVisitor) Visit(n ast.Node) ast.Visitor { + if tp, ok := n.(*ast.TypeSpec); ok { + if _, ok := tp.Type.(*ast.StructType); ok { + ast.Walk(v.visitor(n), n) + return nil // This will ensure this struct is no longer traversed + } + } + return v +} diff --git a/go.mod b/go.mod index 2961bbc..57a6e5f 100644 --- a/go.mod +++ b/go.mod @@ -3,7 +3,8 @@ module github.com/unistack-org/protoc-gen-micro/v3 go 1.16 require ( - github.com/unistack-org/micro-proto v0.0.2-0.20210227213711-77c7563bd01e + github.com/fatih/structtag v1.2.0 + github.com/unistack-org/micro-proto v0.0.2 golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 // indirect google.golang.org/protobuf v1.26.0 ) diff --git a/go.sum b/go.sum index 8ad21b1..de3cb5f 100644 --- a/go.sum +++ b/go.sum @@ -4,6 +4,8 @@ github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw= github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4= github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c= +github.com/fatih/structtag v1.2.0 h1:/OdNE99OxoI/PqaW/SuSK9uxxT3f/tcSZgon/ssNSx4= +github.com/fatih/structtag v1.2.0/go.mod h1:mBJUNpUnHmRKrKlQQlmCrh5PuhftFbNv8Ys4/aAZl94= github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q= github.com/golang/mock v1.1.1/go.mod h1:oTYuIxOrZwtPieC+H1uAHpcLFnEyAGVDL/k47Jfbm0A= github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U= @@ -15,17 +17,17 @@ github.com/golang/protobuf v1.4.0-rc.4.0.20200313231945-b860323f09d0/go.mod h1:W github.com/golang/protobuf v1.4.0/go.mod h1:jodUvKwWbYaEsadDk5Fwe5c77LiNKVO9IDvqG2KuDX0= github.com/golang/protobuf v1.4.1/go.mod h1:U8fpvMrcmy5pZrNK1lt4xCsGvpyWQ/VVv6QDs8UjoX8= github.com/golang/protobuf v1.4.3/go.mod h1:oDoupMAO8OvCJWAcko0GGGIgR6R6ocIYbsSw735rRwI= -github.com/golang/protobuf v1.5.0 h1:LUVKkCeviFUMKqHa4tXIIij/lbhnMbP7Fn5wKdKkRh4= github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk= github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M= github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU= github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/google/go-cmp v0.5.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= +github.com/google/go-cmp v0.5.5 h1:Khx7svrCpmxxtHBq5j2mp/xVjsi8hQMfNLvJFAlrGgU= github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE= github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA= -github.com/unistack-org/micro-proto v0.0.2-0.20210227213711-77c7563bd01e h1:hQJ3V0QggeFdU5967wO5v6oWnaK42wUnG4UU4zWcyu4= -github.com/unistack-org/micro-proto v0.0.2-0.20210227213711-77c7563bd01e/go.mod h1:GYO53DWmeldRIo90cAdQx8bLr/WJMxW62W4ja74p1Ac= +github.com/unistack-org/micro-proto v0.0.2 h1:mL1ZPRGPCWJOiMBiJrkk8Y513rPrJmhJWxyLceckAXU= +github.com/unistack-org/micro-proto v0.0.2/go.mod h1:GYO53DWmeldRIo90cAdQx8bLr/WJMxW62W4ja74p1Ac= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/exp v0.0.0-20190121172915-509febef88a4/go.mod h1:CJ0aWSM057203Lf6IL+f9T1iT9GByDxfZKAQTCR3kQA= golang.org/x/lint v0.0.0-20181026193005-c67002cb31c3/go.mod h1:UVdnD1Gm6xHRNCYTkRU2/jEulfH38KcIWyp/GAMgvoE= diff --git a/http.go b/http.go index 2fd99a6..4ddf7c7 100644 --- a/http.go +++ b/http.go @@ -21,7 +21,7 @@ func (g *Generator) httpGenerate(component string, plugin *protogen.Plugin, genC gfile := plugin.NewGeneratedFile(gname, path) gfile.P("// Code generated by protoc-gen-micro") - gfile.P("// source: ", *file.Proto.Name) + gfile.P("// source: ", file.Proto.GetName()) gfile.P("package ", file.GoPackageName) gfile.P() diff --git a/main.go b/main.go index 296df14..db61c56 100644 --- a/main.go +++ b/main.go @@ -13,6 +13,7 @@ var ( flagDebug = flag.Bool("debug", false, "") flagStandalone = flag.Bool("standalone", false, "") flagComponents = flag.String("components", "micro|rpc|http|client|server", "") + flagTagPath = flag.String("tag_path", ".", "") ) func main() { @@ -29,6 +30,7 @@ type Generator struct { components string standalone bool debug bool + tagPath string } func (g *Generator) Generate(plugin *protogen.Plugin) error { @@ -37,6 +39,7 @@ func (g *Generator) Generate(plugin *protogen.Plugin) error { g.standalone = *flagStandalone g.debug = *flagDebug g.components = *flagComponents + g.tagPath = *flagTagPath plugin.SupportedFeatures = uint64(pluginpb.CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL) var genClient bool @@ -81,5 +84,10 @@ func (g *Generator) Generate(plugin *protogen.Plugin) error { } + if err = g.astGenerate(plugin); err != nil { + plugin.Error(err) + return err + } + return nil } diff --git a/micro.go b/micro.go index fa534a3..fc6edee 100644 --- a/micro.go +++ b/micro.go @@ -22,7 +22,7 @@ func (g *Generator) microGenerate(component string, plugin *protogen.Plugin, gen gfile := plugin.NewGeneratedFile(gname, path) gfile.P("// Code generated by protoc-gen-micro") - gfile.P("// source: ", *file.Proto.Name) + gfile.P("// source: ", file.Proto.GetName()) gfile.P("package ", file.GoPackageName) gfile.P() diff --git a/rpc.go b/rpc.go index 75fa385..ae58ad4 100644 --- a/rpc.go +++ b/rpc.go @@ -21,7 +21,7 @@ func (g *Generator) rpcGenerate(component string, plugin *protogen.Plugin, genCl gfile := plugin.NewGeneratedFile(gname, path) gfile.P("// Code generated by protoc-gen-micro") - gfile.P("// source: ", *file.Proto.Name) + gfile.P("// source: ", file.Proto.GetName()) gfile.P("package ", file.GoPackageName) gfile.P()