将项目添加到Go AST后,注释无序

Question

以下测试尝试使用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%s\nGot:\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)
}

Answer 1

我相信我已经开始工作了.如上述评论所述,要求的要点是:

1. 具体设置包括Slash和的缓冲区位置NamePos
2. 用于token.File.AddLine在特定偏移处添加新行(使用第1项中的位置计算)
3. 全局定位源缓冲区token.File.Position(由源缓冲区使用printer.Printer并且token.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%s\nGot:\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),将所有已分配的字节设置为spaces(0x20)也很重要,这样打印机在处理它们时不会抱怨空字节.

Answer 2

我知道这个答案可能有点晚了。但为了其他人的利益，我在以下 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兼容