Go-AST 获取函数属性

Go提供了ast库,用于将源代码转换成抽象语法树,从而为元编程提供了支持。

一、方法实现

package util

import (
	"fmt"
	"go-less/exception"
	"go/ast"
	"go/parser"
	"go/token"
	"io/ioutil"
	"strings"
)

// ASTFunction AST函数对象
type ASTFunction struct {
	Package  string              // 函数所在包
	Name     string              // 函数名字
	Exported bool                // 是否向外公开
	Recv     []map[string]string // 函数接收者
	Params   []map[string]string // 函数参数
	Results  []map[string]string // 函数返回值
}

// Println Debug Only
func (f ASTFunction) Println() {
	fmt.Print("所在包:" + f.Package)
	fmt.Print(",是否公开:", f.Exported)
	fmt.Print(",函数名字:" + f.Name)
	fmt.Print(",函数接收者:")
	fmt.Print(f.Recv)
	fmt.Print(",函数参数:")
	fmt.Print(f.Params)
	fmt.Print(",函数返回值:")
	fmt.Print(f.Results)
	fmt.Println()
}

// 1.数组或者切片类型 			√
// 2.用户定义的类型或基本数据类型 √
// 3.选择表达式				√
// 4.指针表达式				√
// 5.映射类型					√
// 6.函数类型					√
// 7.管道类型					√
// 8.匿名结构体				×
func exprToTypeStringRecursively(expr ast.Expr) string {

	if arr, ok := expr.(*ast.ArrayType); ok {
		if arr.Len == nil {
			return "[]" + exprToTypeStringRecursively(arr.Elt)
		} else if lit, ok := arr.Len.(*ast.BasicLit); ok {
			return "[" + lit.Value + "]" + exprToTypeStringRecursively(arr.Elt)
		} else {
			// TODO 完备性检查
			panic(1)
		}
	}
	if _, ok := expr.(*ast.InterfaceType); ok {
		return "interface{}"
	}
	if indent, ok := expr.(*ast.Ident); ok {
		return indent.Name
	} else if selExpr, ok := expr.(*ast.SelectorExpr); ok {
		return exprToTypeStringRecursively(selExpr.X) + "." + exprToTypeStringRecursively(selExpr.Sel)
	} else if star, ok := expr.(*ast.StarExpr); ok {
		return "*" + exprToTypeStringRecursively(star.X)
	} else if mapType, ok := expr.(*ast.MapType); ok {
		return "map[" + exprToTypeStringRecursively(mapType.Key) + "]" + exprToTypeStringRecursively(mapType.Value)
	} else if funcType, ok := expr.(*ast.FuncType); ok {
		params := parseFieldList(funcType.Params)
		results := parseFieldList(funcType.Results)
		tf := func(data []map[string]string) string {
			ts := make([]string, 0)
			for _, v := range data {
				ts = append(ts, v["Name"]+" "+v["Type"])
			}
			return strings.Join(ts, ",")
		}
		return "func(" + tf(params) + ")" + " (" + tf(results) + ")"
	} else if chanType, ok := expr.(*ast.ChanType); ok {
		if chanType.Dir == ast.SEND {
			return "chan <- " + exprToTypeStringRecursively(chanType.Value)
		} else if chanType.Dir == ast.RECV {
			return "<- chan " + exprToTypeStringRecursively(chanType.Value)
		} else {
			return "chan " + exprToTypeStringRecursively(chanType.Value)
		}
	}
	//ast.StructType	不考虑这个类型
	// TODO 完备性检查
	fmt.Println(expr)
	panic(1)
}

func parseFieldList(fList *ast.FieldList) []map[string]string {
	dst := make([]map[string]string, 0)
	if fList != nil {
		list := fList.List
		for i := 0; i < len(list); i++ {
			names := list[i].Names
			typeStr := exprToTypeStringRecursively(list[i].Type)
			for j := 0; j < len(names); j++ {
				dst = append(dst, map[string]string{
					"Name": names[j].Name,
					"Type": typeStr,
				})
			}
			if len(names) == 0 {
				dst = append(dst, map[string]string{
					"Name": "",
					"Type": typeStr,
				})
			}
		}
	}
	return dst
}
func CreateASTFunctionFromASTNode(node ast.Node, pack string) *ASTFunction {
	fn, ok := node.(*ast.FuncDecl)
	if ok {
		astFunction := ASTFunction{
			Package:  pack,
			Name:     fn.Name.Name,
			Exported: fn.Name.IsExported(),
		}
		astFunction.Params = parseFieldList(fn.Type.Params)
		astFunction.Results = parseFieldList(fn.Type.Results)
		astFunction.Recv = parseFieldList(fn.Recv)
		return &astFunction
	}
	return nil
}
func CreateASTFunctionsFromFile(target string) []*ASTFunction {
	rawData, err := ioutil.ReadFile(target)
	if err != nil {
		panic(exception.NewDefaultException(exception.IOException, err.Error()))
	}
	fileSet := token.NewFileSet()
	file, err := parser.ParseFile(fileSet, "", string(rawData), 0)
	if err != nil {
		panic(exception.NewDefaultException(exception.ASTParseException, err.Error()))
	}
	pack := ""
	ast.Print(fileSet, file)

	functions := make([]*ASTFunction, 0)
	ast.Inspect(file, func(node ast.Node) bool {
		pk, ok := node.(*ast.Ident)
		if ok {
			if pack == "" {
				pack = pk.Name
			}
		}
		fn := CreateASTFunctionFromASTNode(node, pack)
		if fn != nil {
			functions = append(functions, fn)
		}
		return true
	})
	return functions
}


二、测试代码

functions := util.CreateASTFunctionsFromFile("C:\\Program Files\\Go\\src\\go\\ast\\ast.go")
for _, f := range functions {
	f.Println()
}

三、运行结果

所在包:ast,是否公开:true,函数名字:Pos,函数接收者:[map[Name:c Type:*Comment]],函数参数:[],函数返回值:[map[Name: Type:token.Pos]]
所在包:ast,是否公开:true,函数名字:End,函数接收者:[map[Name:c Type:*Comment]],函数参数:[],函数返回值:[map[Name: Type:token.Pos]]
所在包:ast,是否公开:true,函数名字:Pos,函数接收者:[map[Name:g Type:*CommentGroup]],函数参数:[],函数返回值:[map[Name: Type:token.Pos]]
...
上一篇:在Minecraft里直接当场......


下一篇:AST