persist/cmd/pgen/gen.go

158 lines
3.1 KiB
Go

package main
import (
"io"
"log"
"go/ast"
"go/format"
"go/parser"
"go/token"
"golang.org/x/tools/go/ast/astutil"
"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
name string
}
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
}
var (
qexp *regexp.Regexp
)
func init() {
qexp = regexp.MustCompile(`"`)
}
func NewGenerator(ns ...string) *Generator {
g := &Generator{}
g.firstone = true
g.fset = token.NewFileSet()
if len(ns) > 0 {
g.name = ns[0]
} else {
g.name = "main"
}
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, parser.ParseComments)
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(name string) {
if g.nodes == nil {
return
}
// strip quotation marks from imports
name = qexp.ReplaceAllString(name,"")
astutil.AddImport(g.fset, g.nodes, name)
}
func (g *Generator) Save(of io.Writer) {
if g.nodes == nil {
return
}
if g.name != "main" {
g.nodes.Name = ast.NewIdent(g.name)
}
err := format.Node(of,g.fset,g.nodes)
if err != nil {
log.Fatal("Generate error:",err)
}
}