以下测试尝试使用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:
%s
Got:
%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)
}
I believe I have gotten it to work. As stated in my comment above, the main points required are:
Slash
and NamePos
token.File.AddLine
to add new lines at specific offsets (calculated using the positions from item 1)token.File.Position
(used by printer.Printer
and token.File.Addline
don't fail range checks on the source bufferCode:
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:
%s
Got:
%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)
}
Example: http://play.golang.org/p/_q1xh3giHm
For Item (3), it is also important to set all the overallocated bytes to spaces (0x20
), so that the printer doesn't complain about null bytes when processing them.