godog/builder.go

158 строки
2,9 КиБ
Go

package godog
import (
"bytes"
"go/ast"
"go/format"
"go/parser"
"go/token"
"os"
"path/filepath"
"strings"
"text/template"
"golang.org/x/tools/imports"
)
var mainTpl = `package main
import (
"github.com/DATA-DOG/godog"
)
func main() {
suite := &GodogSuite{
steps: make(map[*regexp.Regexp]StepHandler),
}
{{range $c := .Contexts}}
{{$c}}(suite)
{{end}}
suite.Run()
}
`
type builder struct {
files map[string]*ast.File
fset *token.FileSet
Contexts []string
}
func newBuilder() *builder {
return &builder{
files: make(map[string]*ast.File),
fset: token.NewFileSet(),
}
}
func (b *builder) parseFile(path string) error {
f, err := parser.ParseFile(b.fset, path, nil, 0)
if err != nil {
return err
}
b.deleteMainFunc(f)
b.registerSteps(f)
b.deleteImports(f)
b.files[path] = f
return nil
}
func (b *builder) deleteImports(f *ast.File) {
var decls []ast.Decl
for _, d := range f.Decls {
fun, ok := d.(*ast.GenDecl)
if !ok {
decls = append(decls, d)
continue
}
if fun.Tok != token.IMPORT {
decls = append(decls, fun)
}
}
f.Decls = decls
}
func (b *builder) deleteMainFunc(f *ast.File) {
var decls []ast.Decl
for _, d := range f.Decls {
fun, ok := d.(*ast.FuncDecl)
if !ok {
decls = append(decls, d)
continue
}
if fun.Name.Name != "main" {
decls = append(decls, fun)
}
}
f.Decls = decls
}
func (b *builder) registerSteps(f *ast.File) {
for _, d := range f.Decls {
fun, ok := d.(*ast.FuncDecl)
if !ok {
continue
}
for _, param := range fun.Type.Params.List {
ident, ok := param.Type.(*ast.Ident)
if !ok {
continue
}
if ident.Name == "godog.Suite" || f.Name.Name == "godog" && ident.Name == "Suite" {
b.Contexts = append(b.Contexts, fun.Name.Name)
}
}
}
}
func (b *builder) merge() (*ast.File, error) {
var buf bytes.Buffer
t := template.Must(template.New("main").Parse(mainTpl))
if err := t.Execute(&buf, b); err != nil {
return nil, err
}
f, err := parser.ParseFile(b.fset, "", &buf, 0)
if err != nil {
return nil, err
}
b.deleteImports(f)
b.files["main.go"] = f
pkg, _ := ast.NewPackage(b.fset, b.files, nil, nil)
pkg.Name = "main"
return ast.MergePackageFiles(pkg, ast.FilterImportDuplicates), nil
}
// Build creates a runnable godog executable
// from current package go files
func Build() ([]byte, error) {
b := newBuilder()
err := filepath.Walk(".", func(path string, file os.FileInfo, err error) error {
if file.IsDir() && file.Name() != "." {
return filepath.SkipDir
}
if err == nil && strings.HasSuffix(path, ".go") {
if err := b.parseFile(path); err != nil {
return err
}
}
return err
})
if err != nil {
return nil, err
}
merged, err := b.merge()
if err != nil {
return nil, err
}
var buf bytes.Buffer
if err := format.Node(&buf, b.fset, merged); err != nil {
return nil, err
}
return imports.Process("", buf.Bytes(), nil)
}