godog/pkg/formatters/ast/ast.go

package ast

import (
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"os"
	"strings"

	"errors"
)

const (
	INIT_TEST_GO_FNAME = "init_test.go"
)

type ASTer struct {
	pkg               string
	pkg_test_go_fname string
	init_test_fset    *token.FileSet
	init_test_node    *ast.File
	pkg_test_fset     *token.FileSet
	pkg_test_node     *ast.File
}

func NewASTer() (*ASTer, error) {
	a := &ASTer{}

	pkg, err := получитьИмяGoПакетаВЭтойДире()
	if err != nil {
		return nil, err
	}
	a.pkg = pkg

	err = a.найтиТестовыйФайл()
	if err != nil {
		return nil, err
	}

	return a, nil
}

func (a *ASTer) ДобавитьШаг(шаг, f, ps string) error {
	return a.добавитьШаг(шаг, f, ps)
}

func (a *ASTer) найтиТестовыйФайл() error {
	a.сгенеритьИмяФайла()
	_, err := os.Stat(INIT_TEST_GO_FNAME)
	if err != nil {
		return err
	}

	_, err = os.Stat(a.pkg_test_go_fname)
	return err
}
func (a *ASTer) сгенеритьИмяФайла() {
	a.pkg_test_go_fname = a.pkg + "_test.go"
}

func (a *ASTer) добавитьШаг(шаг, f, ps string) error {
	err := a.добавитьШагвInitФайл(шаг, f)
	if err != nil {
		return err
	}

	err = a.добавитьФункциювТестовыйФайл(f, ps)
	if err != nil {
		return err
	}

	return nil
}
func (a *ASTer) добавитьШагвInitФайл(шаг, f string) error {
	err := a.спарситьInitФайлиДобавитьШаг(шаг, f)
	if err != nil {
		return err
	}
	err = a.перезаписатьInitФайл()
	if err != nil {
		return err
	}
	return nil
}
func (a *ASTer) добавитьФункциювТестовыйФайл(f, ps string) error {
	err := a.спарситьФайлиДобавитьФункцию(f, ps)
	if err != nil {
		return err
	}
	err = a.перезаписатьФайл()
	if err != nil {
		return err
	}
	return nil
}

func (a *ASTer) спарситьInitФайлиДобавитьШаг(шаг, func_name string) error {
	err := a.спарситьInitФайл()
	if err != nil {
		return err
	}

	f := a.создатьШаг(шаг, func_name)
	a.добавитьШагвИнициализаторЕслиНету(f)
	return nil
}
func (a *ASTer) спарситьФайлиДобавитьФункцию(func_name, ps string) error {
	err := a.спарситьФайл()
	if err != nil {
		return err
	}

	f := a.создатьФункцию(func_name, ps)
	a.добавитьФункцию(f)
	return nil
}

func (a *ASTer) добавитьШагвИнициализаторЕслиНету(step *ast.ExprStmt) {
	for _, d := range a.init_test_node.Decls {
		f, ok := d.(*ast.FuncDecl)
		if !ok {
			continue
		}
		if f.Name == nil {
			continue
		}
		if f.Name.Name == "InitializeScenario" {
			a.добавитьШагвФункциюInitializeScenario(f, step)
			return
		}
	}
}
func (a *ASTer) добавитьШагвФункциюInitializeScenario(f *ast.FuncDecl, step *ast.ExprStmt) {
	for i, stmt := range f.Body.List {
		if являетсяЛиШагомсТакойЖеФункцией(stmt, step) {
			return
		}

		if !являетсяЛиШагом(stmt) {
			n := append([]ast.Stmt{step}, f.Body.List[i:]...)
			f.Body.List = append(f.Body.List[0:i], n...)
			return
		}
	}
}
func (a *ASTer) добавитьФункцию(fun *ast.FuncDecl) {
	for _, d := range a.pkg_test_node.Decls {
		f, ok := d.(*ast.FuncDecl)
		if !ok {
			continue
		}
		if f.Name == nil {
			continue
		}
		if f.Name.Name == fun.Name.Name {
			return
		}
	}

	a.pkg_test_node.Decls = append(a.pkg_test_node.Decls, fun)
}

func (a *ASTer) спарситьInitФайл() error {
	fset := token.NewFileSet()
	node, err := parser.ParseFile(fset, INIT_TEST_GO_FNAME, nil, parser.AllErrors|parser.ParseComments)
	a.init_test_fset = fset
	a.init_test_node = node
	return err
}
func (a *ASTer) спарситьФайл() error {
	fset := token.NewFileSet()
	node, err := parser.ParseFile(fset, a.pkg_test_go_fname, nil, parser.AllErrors|parser.ParseComments)
	a.pkg_test_fset = fset
	a.pkg_test_node = node
	return err
}

func (a *ASTer) создатьШаг(шаг, func_name string) *ast.ExprStmt {
	st := &ast.ExprStmt{
		X: &ast.CallExpr{
			Fun: &ast.SelectorExpr{
				X:   ast.NewIdent("ctx"),
				Sel: ast.NewIdent("Step"),
			},
			Args: []ast.Expr{
				&ast.BasicLit{
					Kind:  token.STRING,
					Value: шаг,
				},
				ast.NewIdent(func_name),
			},
		},
	}

	return st
}
func (a *ASTer) создатьФункцию(func_name, ps string) *ast.FuncDecl {
	f := &ast.FuncDecl{
		Name: ast.NewIdent(func_name),
		Type: &ast.FuncType{
			Params: &ast.FieldList{},
		},
		Body: &ast.BlockStmt{},
	}

	a.добавитьПараметрыФункции(f, ps)

	return f
}
func (a *ASTer) добавитьПараметрыФункции(f *ast.FuncDecl, ps string) {
	params := a.сконвертитьСтрокувПараметры(ps)

	var names []*ast.Ident

	for _, p := range params {
		switch len(p) {
		case 0:
			continue
		case 1:
			names = append(names, ast.NewIdent(p[0]))
		case 2:
			names = append(names, ast.NewIdent(p[0]))
			field := &ast.Field{
				Names: names,
				Type:  ast.NewIdent(p[1]),
			}
			f.Type.Params.List = append(f.Type.Params.List, field)
			names = nil
		}
	}
}
func (a *ASTer) сконвертитьСтрокувПараметры(in string) (ret Параметры) {
	params := strings.Split(in, ",")

	for _, p := range params {
		ps := strings.Fields(p)
		параметр := Параметр{}

		switch len(ps) {
		case 0:
			continue
		case 1:
			параметр = append(параметр, ps[0])
		case 2:
			параметр = append(параметр, ps[0], ps[1])
		}
		ret = append(ret, параметр)
	}

	return
}

func (a *ASTer) перезаписатьInitФайл() error {
	err := переименоватьФайлСоВременем(INIT_TEST_GO_FNAME)
	if err != nil {
		return errors.Join(err, errors.New("cant backup orig file"))
	}

	f, err := os.Create(INIT_TEST_GO_FNAME)
	if err != nil {
		return errors.Join(err, errors.New("cant rewrite orig file"))
	}
	defer f.Close()

	return format.Node(f, a.init_test_fset, a.init_test_node)
}
func (a *ASTer) перезаписатьФайл() error {
	err := переименоватьФайлСоВременем(a.pkg_test_go_fname)
	if err != nil {
		return errors.Join(err, errors.New("cant backup orig file"))
	}

	f, err := os.Create(a.pkg_test_go_fname)
	if err != nil {
		return errors.Join(err, errors.New("cant rewrite orig file"))
	}
	defer f.Close()

	return format.Node(f, a.pkg_test_fset, a.pkg_test_node)
}