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 }