Go AST

2023-09-06 ⏳1.1分钟(0.4千字)

最近我司在拆分微服务，在迁移过程中重复工作较多，例如：

于是，一个idea诞生了，这就是func2pb.

而本文意在阐述如何解析Go AST语法树去完成某些特定任务。

获取 AST 文件

我们使用go/parser和go/token两个标准库.

// 文件路径
var path string

data, err := ioutil.ReadFile(path)
if err != nil {
    return
}

fset := token.NewFileSet()
fast, err := parser.ParseFile(fset, file, string(data), parser.ParseComments)

至此我们可以获取一个go/ast中的file结构体

type File struct {
  Doc        *CommentGroup   // associated documentation; or nil
  Package    token.Pos       // position of "package" keyword
  Name       *Ident          // package name
  Decls      []Decl          // top-level declarations; or nil
  Scope      *Scope          // package scope (this file only)
  Imports    []*ImportSpec   // imports in this file
  Unresolved []*Ident        // unresolved identifiers in this file
  Comments   []*CommentGroup // list of all comments in the source file
}

获取 AST 中的函数

看到上述的结构体可能会很困惑，不应该有一个字段是Functions或者Structs么？其实所有的信息都在Decls中。我们可以通过断言fn, isFunc := decl.(*ast.FuncDecl); isFunc获取ast.FuncDecl。

// A FuncDecl node represents a function declaration.
FuncDecl struct {
    Doc  *CommentGroup // associated documentation; or nil
    Recv *FieldList    // receiver (methods); or nil (functions)
    Name *Ident        // function/method name
    Type *FuncType     // function signature: type and value parameters, results, and position of "func" keyword
    Body *BlockStmt    // function body; or nil for external (non-Go) function
}

从这个结构体我们就可以很清晰的获取：

获取 AST 中的结构体

而获取结构体相对而言会更复杂。我们需要先将Decls断言genDecl, isGenDecl := decl.(*ast.GenDecl);，获取GenDecl

GenDecl struct {
    Doc    *CommentGroup // associated documentation; or nil
    TokPos token.Pos     // position of Tok
    Tok    token.Token   // IMPORT, CONST, TYPE, or VAR
    Lparen token.Pos     // position of '(', if any
    Specs  []Spec
    Rparen token.Pos // position of ')', if any
}

再GenDecl.Specs断言typeSpec, isTypeSpec := spec.(*ast.TypeSpec);，获取TypeSpec

// A TypeSpec node represents a type declaration (TypeSpec production).
TypeSpec struct {
    Doc        *CommentGroup // associated documentation; or nil
    Name       *Ident        // type name
    TypeParams *FieldList    // type parameters; or nil
    Assign     token.Pos     // position of '=', if any
    Type       Expr          // *Ident, *ParenExpr, *SelectorExpr, *StarExpr, or any of the *XxxTypes
    Comment    *CommentGroup // line comments; or nil
}

最终断言TypeSpec为structType, isStruct := typeSpec.Type.(*ast.StructType);,获取StructType

// A StructType node represents a struct type.
StructType struct {
    Struct     token.Pos  // position of "struct" keyword
    Fields     *FieldList // list of field declarations
    Incomplete bool       // true if (source) fields are missing in the Fields list
}

杂项

关于Go AST的解析强烈建议使用ChatGPT辅助生成。里面逻辑非常复杂。当我们获取到我们想要的内容后，我们可以运用一些辅助工具去生成我们的小工具。

生成命令行

这里为选用的是github.com/urfave/cli.具体不在赘述

生成模版文件

我们使用text/template去生成模版文件,同时还可以注入一些函数template.FuncMap比如转驼峰/字符串替换等等。例子如：

const tmplPB = `syntax = "proto3";

// TODO fill it
package xxx;

service {{ .Sn }} { {{ range .Funcs }}
{{ if commentnotempty .Comment }}       {{ comment .Comment }}{{ end }} rpc {{ .Name }}({{ .Name }}Req) returns ({{ .Name }}Resp);{{ end }}
}

{{ range .Funcs }}
message {{ .Name }}Req {{ "{" }}{{ range $index, $element := .In }}
        {{ replace .Typ }} {{ rename .Name }} = {{$index | add 1}};{{ end }}
}

message {{ .Name }}Resp {
        int32 code = 1;
        string msg = 2;
        {{ .Name }}Data data = 3;
}

message {{ .Name }}Data {{ "{" }}{{ range $index, $element := .Out }}
        {{ replace .Typ }} {{ rename .Name }} = {{$index | add 1}};{{ end }}
}
{{ end }}

{{ range .Structs }}
{{ if commentnotempty .Comment }}{{ comment .Comment }}{{ end }}message {{ .Name }} { {{ range $index, $element := .Field }}
{{ if commentnotempty .Comment }}       {{ comment .Comment }}{{ end }} {{ replace .Typ }} {{ rename .Name }} = {{$index | add 1}};{{ end }}
}
{{ end }}
`

总结

用一些工具思维去解决一些定势问题，会有不少收获。