gobuf/gobuf.go
2024-03-11 00:46:06 +08:00

141 lines
3.2 KiB
Go

package gobuf
import (
"bytes"
"fmt"
"io"
"os"
"github.com/emicklei/proto"
"gitter.top/sync/proto-contrib/pkg/protofmt"
)
type Parser struct {
filename string
rawData io.Reader
proto *proto.Proto
}
func NewParser(file string) (*Parser, error) {
reader, err := os.Open(file)
if err != nil {
return nil, err
}
defer reader.Close()
parser := proto.NewParser(reader)
definition, err := parser.Parse()
if err != nil {
return nil, err
}
return &Parser{
filename: file,
rawData: reader,
proto: definition,
}, nil
}
func (parser *Parser) ExistService(serviceName string) bool {
var result bool
proto.Walk(parser.proto, proto.WithService(func(service *proto.Service) {
if service.Name == serviceName {
result = true
}
}))
return result
}
func (parser *Parser) ExistMessage(messageName string) bool {
var result bool
proto.Walk(parser.proto, proto.WithMessage(func(message *proto.Message) {
if message.Name == messageName {
result = true
}
}))
return result
}
func (parser *Parser) ExistRPC(serviceName, rpcName string) bool {
var result bool
proto.Walk(parser.proto, proto.WithService(func(service *proto.Service) {
if service.Name != serviceName {
return
}
for _, element := range service.Elements {
if rpc, ok := element.(*proto.RPC); ok {
if rpc.Name == rpcName {
result = true
return
}
}
}
}))
return result
}
func (parser *Parser) AddRPC(serviceName, rpcName string) error {
proto.Walk(parser.proto, proto.WithService(func(service *proto.Service) {
if service.Name != serviceName {
return
}
service.Elements = append(service.Elements, &proto.RPC{
Comment: &proto.Comment{
Lines: []string{
" @desc: 描述",
" @author: 作者",
" @method: GET",
" @api: /" + calm2Case(rpcName),
},
},
Name: rpcName,
RequestType: toUpperCamlCase(serviceName) + toUpperCamlCase(rpcName) + "Req",
ReturnsType: toUpperCamlCase(serviceName) + toUpperCamlCase(rpcName) + "Resp",
Parent: service,
})
}))
parser.proto.Elements = append(parser.proto.Elements, &proto.Message{
Name: toUpperCamlCase(serviceName) + toUpperCamlCase(rpcName) + "Req",
Parent: parser.proto,
})
parser.proto.Elements = append(parser.proto.Elements, &proto.Message{
Name: toUpperCamlCase(serviceName) + toUpperCamlCase(rpcName) + "Resp",
Parent: parser.proto,
})
return parser.writeSync()
}
func (parser *Parser) AddService(serviceName string) error {
if parser.ExistService(serviceName) {
return fmt.Errorf("service name exist")
}
parser.proto.Elements = append(parser.proto.Elements, &proto.Service{
Comment: &proto.Comment{
Lines: []string{
" @route_group: true",
" @base_url: /v1/" + calm2Case(serviceName),
" @gen_to: ./services/controller/v1/" + calm2Case(serviceName) + "_controller.go",
},
},
Name: serviceName,
Parent: parser.proto,
})
return parser.writeSync()
}
func (parser *Parser) writeSync() error {
var buf = new(bytes.Buffer)
protofmt.NewFormatter(buf, " ").Format(parser.proto) // 1 tab
// write back to input
if err := os.WriteFile(parser.filename, buf.Bytes(), os.ModePerm); err != nil {
return err
}
return nil
}