persist/cmd/pgen/gen.go

package main

import (
	"path/filepath"
	"log"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"golang.org/x/tools/go/ast/astutil"
	"os"
	"regexp"
	"strings"
)

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() {
	fs,err := filepath.Glob("pgen*.go")
	if err != nil {
		log.Fatal("Glob error")
	}
	for _,f := range(fs) {
		err := os.Remove(f)
		if err != nil && os.IsExist(err) {
			log.Fatal("Removing ",f,err)
		}
	}
	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() {
	if g.nodes == nil {
		return
	}
	var of *os.File
	var err error
	if g.name == "main" {
		of, err = os.Create("pgen.go")
	} else {
		oname := strings.Join([]string{"pgen_",g.name,".go"},"")
		of, err = os.Create(oname)
		g.nodes.Name = ast.NewIdent(g.name)
	}
	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)
	}
}