add tag support
Signed-off-by: Vasiliy Tolstov <v.tolstov@unistack.org>
This commit is contained in:
parent
f69088bd27
commit
d41fa1a64f
153
ast.go
Normal file
153
ast.go
Normal file
@ -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
|
||||
}
|
||||
3
go.mod
3
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
|
||||
)
|
||||
|
||||
8
go.sum
8
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=
|
||||
|
||||
2
http.go
2
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()
|
||||
|
||||
|
||||
8
main.go
8
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
|
||||
}
|
||||
|
||||
2
micro.go
2
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()
|
||||
|
||||
|
||||
2
rpc.go
2
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()
|
||||
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user