You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
go-zero/tools/goctl/rpc/parser/pbast.go

481 lines
11 KiB
Go

package parser
import (
"errors"
"fmt"
"go/ast"
"go/parser"
"go/token"
"sort"
"strings"
"github.com/tal-tech/go-zero/core/lang"
sx "github.com/tal-tech/go-zero/core/stringx"
"github.com/tal-tech/go-zero/tools/goctl/util"
"github.com/tal-tech/go-zero/tools/goctl/util/console"
"github.com/tal-tech/go-zero/tools/goctl/util/stringx"
)
const (
flagStar = "*"
suffixServer = "Server"
referenceContext = "context."
unknownPrefix = "XXX_"
ignoreJsonTagExpression = `json:"-"`
)
var (
errorParseError = errors.New("pb parse error")
typeTemplate = `type (
{{.types}}
)`
structTemplate = `{{if .type}}type {{end}}{{.name}} struct {
{{.fields}}
}`
fieldTemplate = `{{if .hasDoc}}{{.doc}}
{{end}}{{.name}} {{.type}} {{.tag}}{{if .hasComment}}{{.comment}}{{end}}`
objectM = make(map[string]*Struct)
)
type (
astParser struct {
golang []byte
filterStruct map[string]lang.PlaceholderType
console.Console
fileSet *token.FileSet
}
Field struct {
Name stringx.String
TypeName string
JsonTag string
Document []string
Comment []string
}
Struct struct {
Name stringx.String
Document []string
Comment []string
Field []*Field
}
Func struct {
Name stringx.String
InType string
InTypeName string // remove *Context,such as LoginRequest、UserRequest
OutTypeName string // remove *Context
OutType string
Document []string
}
RpcService struct {
Name stringx.String
Funcs []*Func
}
// parsing for rpc
PbAst struct {
Package string
// external reference
Imports map[string]string
Strcuts map[string]*Struct
// rpc server's functions,not all functions
Service []*RpcService
}
)
func NewAstParser(golang []byte, filterStruct map[string]lang.PlaceholderType, log console.Console) *astParser {
return &astParser{
golang: golang,
filterStruct: filterStruct,
Console: log,
fileSet: token.NewFileSet(),
}
}
func (a *astParser) Parse() (*PbAst, error) {
fSet := a.fileSet
f, err := parser.ParseFile(fSet, "", a.golang, parser.ParseComments)
if err != nil {
return nil, err
}
commentMap := ast.NewCommentMap(fSet, f, f.Comments)
f.Comments = commentMap.Filter(f).Comments()
var pbAst PbAst
pbAst.Package = a.mustGetIndentName(f.Name)
imports := make(map[string]string)
for _, item := range f.Imports {
if item == nil {
continue
}
if item.Path == nil {
continue
}
key := a.mustGetIndentName(item.Name)
value := item.Path.Value
imports[key] = value
}
structs, funcs := a.mustScope(f.Scope)
pbAst.Imports = imports
pbAst.Strcuts = structs
pbAst.Service = funcs
return &pbAst, nil
}
func (a *astParser) mustScope(scope *ast.Scope) (map[string]*Struct, []*RpcService) {
if scope == nil {
return nil, nil
}
objects := scope.Objects
structs := make(map[string]*Struct)
serviceList := make([]*RpcService, 0)
for name, obj := range objects {
decl := obj.Decl
if decl == nil {
continue
}
typeSpec, ok := decl.(*ast.TypeSpec)
if !ok {
continue
}
tp := typeSpec.Type
switch v := tp.(type) {
case *ast.StructType:
st, err := a.parseObject(name, v)
a.Must(err)
structs[st.Name.Lower()] = st
case *ast.InterfaceType:
if !strings.HasSuffix(name, suffixServer) {
continue
}
list := a.mustServerFunctions(v)
serviceList = append(serviceList, &RpcService{
Name: stringx.From(strings.TrimSuffix(name, suffixServer)),
Funcs: list,
})
}
}
targetStruct := make(map[string]*Struct)
for st := range a.filterStruct {
lower := strings.ToLower(st)
targetStruct[lower] = structs[lower]
}
return targetStruct, serviceList
}
func (a *astParser) mustServerFunctions(v *ast.InterfaceType) []*Func {
funcs := make([]*Func, 0)
methodObject := v.Methods
if methodObject == nil {
return nil
}
for _, method := range methodObject.List {
var item Func
name := a.mustGetIndentName(method.Names[0])
doc := a.parseCommentOrDoc(method.Doc)
item.Name = stringx.From(name)
item.Document = doc
types := method.Type
if types == nil {
funcs = append(funcs, &item)
continue
}
v, ok := types.(*ast.FuncType)
if !ok {
continue
}
params := v.Params
if params != nil {
inList, err := a.parseFields(params.List, true)
a.Must(err)
for _, data := range inList {
if strings.HasPrefix(data.TypeName, referenceContext) {
continue
}
// currently,does not support external references
item.InTypeName = data.TypeName
item.InType = strings.TrimPrefix(data.TypeName, flagStar)
break
}
}
results := v.Results
if results != nil {
outList, err := a.parseFields(results.List, true)
a.Must(err)
for _, data := range outList {
if strings.HasPrefix(data.TypeName, referenceContext) {
continue
}
// currently,does not support external references
item.OutTypeName = data.TypeName
item.OutType = strings.TrimPrefix(data.TypeName, flagStar)
break
}
}
funcs = append(funcs, &item)
}
return funcs
}
func (a *astParser) parseObject(structName string, tp *ast.StructType) (*Struct, error) {
if data, ok := objectM[structName]; ok {
return data, nil
}
var st Struct
st.Name = stringx.From(structName)
if tp == nil {
return &st, nil
}
fields := tp.Fields
if fields == nil {
objectM[structName] = &st
return &st, nil
}
fieldList := fields.List
members, err := a.parseFields(fieldList, false)
if err != nil {
return nil, err
}
for _, m := range members {
var field Field
field.Name = m.Name
field.TypeName = m.TypeName
field.JsonTag = m.JsonTag
field.Document = m.Document
field.Comment = m.Comment
st.Field = append(st.Field, &field)
}
objectM[structName] = &st
return &st, nil
}
func (a *astParser) parseFields(fields []*ast.Field, onlyType bool) ([]*Field, error) {
ret := make([]*Field, 0)
for _, field := range fields {
var item Field
tag := a.parseTag(field.Tag)
if tag == "" && !onlyType {
continue
}
if tag == ignoreJsonTagExpression {
continue
}
item.JsonTag = tag
name := a.parseName(field.Names)
if strings.HasPrefix(name, unknownPrefix) {
continue
}
item.Name = stringx.From(name)
typeName, err := a.parseType(field.Type)
if err != nil {
return nil, err
}
item.TypeName = typeName
if onlyType {
ret = append(ret, &item)
continue
}
docs := a.parseCommentOrDoc(field.Doc)
comments := a.parseCommentOrDoc(field.Comment)
item.Document = docs
item.Comment = comments
isInline := name == ""
if isInline {
return nil, a.wrapError(field.Pos(), "unexpected inline type:%s", name)
}
ret = append(ret, &item)
}
return ret, nil
}
func (a *astParser) parseTag(basicLit *ast.BasicLit) string {
if basicLit == nil {
return ""
}
value := basicLit.Value
splits := strings.Split(value, " ")
if len(splits) == 1 {
return fmt.Sprintf("`%s`", strings.ReplaceAll(splits[0], "`", ""))
} else {
return fmt.Sprintf("`%s`", strings.ReplaceAll(splits[1], "`", ""))
}
}
// returns
// resp1:type's string expression,like int、string、[]int64、map[string]User、*User
// resp2:error
func (a *astParser) parseType(expr ast.Expr) (string, error) {
if expr == nil {
return "", errorParseError
}
switch v := expr.(type) {
case *ast.StarExpr:
stringExpr, err := a.parseType(v.X)
if err != nil {
return "", err
}
e := fmt.Sprintf("*%s", stringExpr)
return e, nil
case *ast.Ident:
return a.mustGetIndentName(v), nil
case *ast.MapType:
keyStringExpr, err := a.parseType(v.Key)
if err != nil {
return "", err
}
valueStringExpr, err := a.parseType(v.Value)
if err != nil {
return "", err
}
e := fmt.Sprintf("map[%s]%s", keyStringExpr, valueStringExpr)
return e, nil
case *ast.ArrayType:
stringExpr, err := a.parseType(v.Elt)
if err != nil {
return "", err
}
e := fmt.Sprintf("[]%s", stringExpr)
return e, nil
case *ast.InterfaceType:
return "interface{}", nil
case *ast.SelectorExpr:
join := make([]string, 0)
xIdent, ok := v.X.(*ast.Ident)
xIndentName := a.mustGetIndentName(xIdent)
if ok {
join = append(join, xIndentName)
}
sel := v.Sel
join = append(join, a.mustGetIndentName(sel))
return strings.Join(join, "."), nil
case *ast.ChanType:
return "", a.wrapError(v.Pos(), "unexpected type 'chan'")
case *ast.FuncType:
return "", a.wrapError(v.Pos(), "unexpected type 'func'")
case *ast.StructType:
return "", a.wrapError(v.Pos(), "unexpected inline struct type")
default:
return "", a.wrapError(v.Pos(), "unexpected type '%v'", v)
}
}
func (a *astParser) parseName(names []*ast.Ident) string {
if len(names) == 0 {
return ""
}
name := names[0]
return a.mustGetIndentName(name)
}
func (a *astParser) parseCommentOrDoc(cg *ast.CommentGroup) []string {
if cg == nil {
return nil
}
comments := make([]string, 0)
for _, comment := range cg.List {
if comment == nil {
continue
}
text := strings.TrimSpace(comment.Text)
if text == "" {
continue
}
comments = append(comments, text)
}
return comments
}
func (a *astParser) mustGetIndentName(ident *ast.Ident) string {
if ident == nil {
return ""
}
return ident.Name
}
func (a *astParser) wrapError(pos token.Pos, format string, arg ...interface{}) error {
file := a.fileSet.Position(pos)
return fmt.Errorf("line %v: %s", file.Line, fmt.Sprintf(format, arg...))
}
func (a *PbAst) GenTypesCode() (string, error) {
types := make([]string, 0)
sts := make([]*Struct, 0)
for _, item := range a.Strcuts {
sts = append(sts, item)
}
sort.Slice(sts, func(i, j int) bool {
return sts[i].Name.Source() < sts[j].Name.Source()
})
for _, s := range sts {
structCode, err := s.genCode(false)
if err != nil {
return "", err
}
if structCode == "" {
continue
}
types = append(types, structCode)
}
buffer, err := util.With("type").Parse(typeTemplate).Execute(map[string]interface{}{
"types": strings.Join(types, "\n\n"),
})
if err != nil {
return "", err
}
return buffer.String(), nil
}
func (s *Struct) genCode(containsTypeStatement bool) (string, error) {
fields := make([]string, 0)
for _, f := range s.Field {
var comment, doc string
if len(f.Comment) > 0 {
comment = f.Comment[0]
}
doc = strings.Join(f.Document, "\n")
buffer, err := util.With(sx.Rand()).Parse(fieldTemplate).Execute(map[string]interface{}{
"name": f.Name.Title(),
"type": f.TypeName,
"tag": f.JsonTag,
"hasDoc": len(f.Document) > 0,
"doc": doc,
"hasComment": len(f.Comment) > 0,
"comment": comment,
})
if err != nil {
return "", err
}
fields = append(fields, buffer.String())
}
buffer, err := util.With("struct").Parse(structTemplate).Execute(map[string]interface{}{
"type": containsTypeStatement,
"name": s.Name.Title(),
"fields": strings.Join(fields, "\n"),
})
if err != nil {
return "", err
}
return buffer.String(), nil
}