155 lines
3.1 KiB
Go
155 lines
3.1 KiB
Go
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
|
|
}
|
|
|
|
type Generator struct{
|
|
firstone bool
|
|
fset *token.FileSet
|
|
fspec, tspec, nspec subspec
|
|
nspecs []subspec
|
|
nodes *ast.File
|
|
}
|
|
|
|
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 (g *Generator) 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,g.nspecs)
|
|
subs(&n.Obj.Name,g.nspecs)
|
|
}
|
|
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,g.nspecs)
|
|
}
|
|
}
|
|
case *ast.TypeAssertExpr:
|
|
nt, ok := n.Type.(*ast.Ident); if !ok { return g }
|
|
subs(&nt.Name,g.nspecs)
|
|
case *ast.Ellipsis:
|
|
nt, ok := n.Elt.(*ast.Ident); if !ok { return g }
|
|
subs(&nt.Name,g.nspecs)
|
|
case *ast.Field:
|
|
nt, ok := n.Type.(*ast.Ident); if !ok { return g }
|
|
subs(&nt.Name,g.nspecs)
|
|
case *ast.StarExpr:
|
|
nt, ok := n.X.(*ast.Ident); if !ok { return g }
|
|
subs(&nt.Name,g.nspecs)
|
|
case *ast.ValueSpec:
|
|
nt, ok := n.Type.(*ast.Ident); if !ok { return g }
|
|
subs(&nt.Name,g.nspecs)
|
|
case *ast.ArrayType:
|
|
nt, ok := n.Elt.(*ast.Ident); if !ok { return g }
|
|
subs(&nt.Name,g.nspecs)
|
|
case *ast.MapType:
|
|
nt, ok := n.Key.(*ast.Ident); if ok {
|
|
subs(&nt.Name,g.nspecs)
|
|
}
|
|
nt, ok = n.Value.(*ast.Ident); if ok {
|
|
subs(&nt.Name,g.nspecs)
|
|
}
|
|
case *ast.FuncDecl:
|
|
sub(&n.Name.Name,g.fspec)
|
|
//ast.Print(g.fset,n)
|
|
}
|
|
return g
|
|
}
|
|
|
|
func init() {
|
|
err := os.Remove("pgen.go")
|
|
if err != nil && os.IsExist(err) {
|
|
log.Fatal("Removing pgen.go:",err)
|
|
}
|
|
}
|
|
|
|
func NewGenerator() *Generator {
|
|
g := &Generator{}
|
|
g.firstone = true
|
|
g.fset = token.NewFileSet()
|
|
return g
|
|
}
|
|
|
|
func (g *Generator) Add(newName, typName, typ string) {
|
|
g.fspec = subspec{regexp.MustCompile("New"), newName}
|
|
g.nspecs = []subspec{
|
|
subspec{regexp.MustCompile("_N"), typName},
|
|
subspec{regexp.MustCompile("_T"), typ}}
|
|
|
|
f, err := parser.ParseFile(g.fset, "", template, 0)
|
|
if err != nil {
|
|
log.Fatal("Parsing persist/template.T:",err)
|
|
}
|
|
ast.Walk(g,f)
|
|
|
|
if g.firstone {
|
|
g.nodes = f
|
|
} else {
|
|
for _,decl := range(f.Decls) {
|
|
imp, ok := decl.(*ast.GenDecl); if ok {
|
|
if imp.Tok == token.IMPORT {
|
|
continue // skip imports
|
|
}
|
|
}
|
|
g.nodes.Decls = append(g.nodes.Decls,decl)
|
|
}
|
|
}
|
|
g.firstone = false
|
|
}
|
|
|
|
func (g *Generator) Import(imports map[string]bool) {
|
|
if g.nodes == nil {
|
|
return
|
|
}
|
|
|
|
// strip quotation marks from imports
|
|
qexp := regexp.MustCompile(`"`)
|
|
for v := range imports {
|
|
v = qexp.ReplaceAllString(v,"")
|
|
astutil.AddImport(g.fset, g.nodes, v)
|
|
}
|
|
}
|
|
|
|
func (g *Generator) Save() {
|
|
if g.nodes == nil {
|
|
return
|
|
}
|
|
of, err := os.Create("pgen.go")
|
|
if err != nil {
|
|
log.Fatal("Cannot open pgen.go")
|
|
}
|
|
|
|
err = format.Node(of,g.fset,g.nodes)
|
|
if err != nil {
|
|
log.Fatal("Generate error:",err)
|
|
}
|
|
}
|
|
|