package main import ( "log" "go/ast" "go/format" "go/parser" "go/token" "golang.org/x/tools/go/ast/astutil" "os" "regexp" ) type subspec struct { r *regexp.Regexp s string } var ( firstone bool gfset *token.FileSet fspec, tspec, nspec subspec ntspecs []subspec nodes *ast.File ) type gvisit struct{} func subs(dst *string, specs []subspec) { for _,spec := range(specs) { sub(dst,spec) } } func sub(dst *string,spec subspec) { if spec.r.MatchString(*dst) { *dst = spec.r.ReplaceAllString(*dst,spec.s) } } func (v *gvisit) Visit(node ast.Node) ast.Visitor { switch n := node.(type) { case *ast.Ident: if n.Obj != nil && n.Obj.Kind == ast.Typ { subs(&n.Name,ntspecs) subs(&n.Obj.Name,ntspecs) } case *ast.TypeSwitchStmt: for _,s := range(n.Body.List) { cs, ok := s.(*ast.CaseClause); if !ok { continue } for _,c := range(cs.List) { ci, ok := c.(*ast.Ident); if !ok { continue } subs(&ci.Name,ntspecs) } } case *ast.TypeAssertExpr: nt, ok := n.Type.(*ast.Ident); if !ok { return v } subs(&nt.Name,ntspecs) case *ast.Ellipsis: nt, ok := n.Elt.(*ast.Ident); if !ok { return v } subs(&nt.Name,ntspecs) case *ast.Field: nt, ok := n.Type.(*ast.Ident); if !ok { return v } subs(&nt.Name,ntspecs) case *ast.StarExpr: nt, ok := n.X.(*ast.Ident); if !ok { return v } subs(&nt.Name,ntspecs) case *ast.TypeSpec: //ast.Print(gfset,n) case *ast.ValueSpec: nt, ok := n.Type.(*ast.Ident); if !ok { return v } subs(&nt.Name,ntspecs) case *ast.ArrayType: nt, ok := n.Elt.(*ast.Ident); if !ok { return v } subs(&nt.Name,ntspecs) case *ast.MapType: nt, ok := n.Key.(*ast.Ident); if ok { subs(&nt.Name,ntspecs) } nt, ok = n.Value.(*ast.Ident); if ok { subs(&nt.Name,ntspecs) } case *ast.FuncDecl: sub(&n.Name.Name,fspec) //ast.Print(gfset,n) } return v } func init() { firstone = true err := os.Remove("pgen.go") if err != nil && os.IsExist(err) { log.Fatal("Removing pgen.go:",err) } } func add(newName, typName, typ string) { fspec = subspec{regexp.MustCompile("New"), newName} nspec = subspec{regexp.MustCompile("_N"), typName} tspec = subspec{regexp.MustCompile("_T"), typ} ntspecs = []subspec{nspec,tspec} gfset = token.NewFileSet() f, err := parser.ParseFile(gfset, "", template, 0) if err != nil { log.Fatal("Parsing persist/template.T:",err) } var v = &gvisit{} ast.Walk(v,f) if firstone { nodes = f } else { for _,decl := range(f.Decls) { imp, ok := decl.(*ast.GenDecl); if ok { if imp.Tok == token.IMPORT { continue // skip imports } } nodes.Decls = append(nodes.Decls,decl) } } firstone = false } func newImportSpec(name string) *ast.ImportSpec { path := &ast.BasicLit{ Value: name } ret := &ast.ImportSpec{ Path: path } return ret } func gen(wantImps []string) { of, err := os.Create("pgen.go") if err != nil { log.Fatal("Cannot open pgen.go") } qexp := regexp.MustCompile(`"`) for _,v := range wantImps { v = qexp.ReplaceAllString(v,"") astutil.AddImport(gfset, nodes, v) } err = format.Node(of,gfset,nodes) if err != nil { log.Fatal("Generate error:",err) } }