添加项目到Go AST后,注释顺序乱了



下面的测试尝试使用AST向结构体添加字段。字段被正确添加,但是注释被乱序添加。我收集的位置可能需要手动指定,但我一直在寻找一个空白的答案。

这是一个失败的测试:http://play.golang.org/p/RID4N30FZK

代码如下:

package generator
import (
    "bytes"
    "fmt"
    "go/ast"
    "go/parser"
    "go/printer"
    "go/token"
    "testing"
)
func TestAst(t *testing.T) {
    source := `package a
// B comment
type B struct {
    // C comment
    C string
}`
    fset := token.NewFileSet()
    file, err := parser.ParseFile(fset, "", []byte(source), parser.ParseComments)
    if err != nil {
        t.Error(err)
    }
    v := &visitor{
        file: file,
    }
    ast.Walk(v, file)
    var output []byte
    buf := bytes.NewBuffer(output)
    if err := printer.Fprint(buf, fset, file); err != nil {
        t.Error(err)
    }
    expected := `package a
// B comment
type B struct {
    // C comment
    C string
    // D comment
    D int
    // E comment
    E float64
}
`
    if buf.String() != expected {
        t.Error(fmt.Sprintf("Test failed. Expected:n%snGot:n%s", expected, buf.String()))
    }
    /*
    actual output = `package a
// B comment
type B struct {
    // C comment
    // D comment
    // E comment
    C   string
    D   int
    E   float64
}
`
    */
}
type visitor struct {
    file *ast.File
}
func (v *visitor) Visit(node ast.Node) (w ast.Visitor) {
    if node == nil {
        return v
    }
    switch n := node.(type) {
    case *ast.GenDecl:
        if n.Tok != token.TYPE {
            break
        }
        ts := n.Specs[0].(*ast.TypeSpec)
        if ts.Name.Name == "B" {
            fields := ts.Type.(*ast.StructType).Fields
            addStructField(fields, v.file, "int", "D", "D comment")
            addStructField(fields, v.file, "float64", "E", "E comment")
        }
    }
    return v
}
func addStructField(fields *ast.FieldList, file *ast.File, typ string, name string, comment string) {
    c := &ast.Comment{Text: fmt.Sprint("// ", comment)}
    cg := &ast.CommentGroup{List: []*ast.Comment{c}}
    f := &ast.Field{
        Doc:   cg,
        Names: []*ast.Ident{ast.NewIdent(name)},
        Type:  ast.NewIdent(typ),
    }
    fields.List = append(fields.List, f)
    file.Comments = append(file.Comments, cg)
}

我相信我已经让它工作了。正如我在上面的评论中所述,需要的要点是:

  1. 具体设置缓冲区位置,包括SlashNamePos
  2. 使用token.File.AddLine在特定偏移量处添加新行(使用第1项的位置计算)
  3. 过度分配源缓冲区,以便token.File.Position(由printer.Printertoken.File.Addline使用)不会在源缓冲区上失败范围检查
代码:

package main
import (
    "bytes"
    "fmt"
    "go/ast"
    "go/parser"
    "go/printer"
    "go/token"
    "testing"
)
func main() {
    tests := []testing.InternalTest{{"TestAst", TestAst}}
    matchAll := func(t string, pat string) (bool, error) { return true, nil }
    testing.Main(matchAll, tests, nil, nil)
}
func TestAst(t *testing.T) {
    source := `package a
// B comment
type B struct {
    // C comment
    C string
}`
    buffer := make([]byte, 1024, 1024)
    for idx,_ := range buffer {
        buffer[idx] = 0x20
    }
    copy(buffer[:], source)
    fset := token.NewFileSet()
    file, err := parser.ParseFile(fset, "", buffer, parser.ParseComments)
    if err != nil {
        t.Error(err)
    }
    v := &visitor{
        file: file,
        fset: fset,
    }
    ast.Walk(v, file)
    var output []byte
    buf := bytes.NewBuffer(output)
    if err := printer.Fprint(buf, fset, file); err != nil {
        t.Error(err)
    }
    expected := `package a
// B comment
type B struct {
    // C comment
    C   string
    // D comment
    D   int
    // E comment
    E   float64
}
`
    if buf.String() != expected {
        t.Error(fmt.Sprintf("Test failed. Expected:n%snGot:n%s", expected, buf.String()))
    }
}
type visitor struct {
    file *ast.File
    fset *token.FileSet
}
func (v *visitor) Visit(node ast.Node) (w ast.Visitor) {
    if node == nil {
        return v
    }
    switch n := node.(type) {
    case *ast.GenDecl:
        if n.Tok != token.TYPE {
            break
        }
        ts := n.Specs[0].(*ast.TypeSpec)
        if ts.Name.Name == "B" {
            fields := ts.Type.(*ast.StructType).Fields
            addStructField(v.fset, fields, v.file, "int", "D", "D comment")
            addStructField(v.fset, fields, v.file, "float64", "E", "E comment")
        }
    }
    return v
}
func addStructField(fset *token.FileSet, fields *ast.FieldList, file *ast.File, typ string, name string, comment string) {
    prevField := fields.List[fields.NumFields()-1] 
    c := &ast.Comment{Text: fmt.Sprint("// ", comment), Slash: prevField.End() + 1}
    cg := &ast.CommentGroup{List: []*ast.Comment{c}}
    o := ast.NewObj(ast.Var, name)
    f := &ast.Field{
        Doc:   cg,
        Names: []*ast.Ident{&ast.Ident{Name: name, Obj: o, NamePos: cg.End() + 1}},
    }
    o.Decl = f
    f.Type = &ast.Ident{Name: typ, NamePos: f.Names[0].End() + 1}
    fset.File(c.End()).AddLine(int(c.End()))
    fset.File(f.End()).AddLine(int(f.End()))
    fields.List = append(fields.List, f)
    file.Comments = append(file.Comments, cg)
}

示例:http://play.golang.org/p/_q1xh3giHm

对于Item (3),将所有超额分配的字节设置为空格(0x20)也很重要,以便打印机在处理它们时不会抱怨空字节。

我知道这个答案可能有点晚了。但是为了其他人的利益,我在以下GitHub问题

中找到了这个库的参考。https://github.com/golang/go/issues/20744

该库名为dst,它可以将go ast转换为dst,反之亦然。

https://github.com/dave/dst

ast中,注释按字节偏移量存储,而不是附加到节点上。Dst通过将注释附加到其各自的节点来解决这个问题,以便重新排列节点不会破坏输出/树。

这个库和广告上说的一样,到目前为止我还没有发现任何问题。

注意:还有一个名为dst/dstutil的子包,它与golang.org/x/tools/go/ast/astutil

兼容

最新更新