added two new functions: AddGlobal() and ReplaceBlock()

gogen/add_global.go

@@ -0,0 +1,73 @@
+package gogen
+
+import (
+	"fmt"
+	"go/ast"
+	"go/token"
+	"strings"
+
+	"github.com/gobuffalo/genny/v2"
+)
+
+// AddGlobal adds global variables.
+func AddGlobal(gf genny.File, globals ...string) (genny.File, error) {
+	pf, err := ParseFile(gf)
+	if err != nil {
+		return gf, err
+	}
+
+	gf = pf.File
+
+	end, isGroup := findLastGlobal(pf.Ast, pf.FileSet, pf.Lines)
+	if end < 0 {
+		return gf, fmt.Errorf("unable to find the position to add the variables")
+	}
+
+	x := []string{}
+	if isGroup {
+		for _, gv := range globals {
+			x = append(x, "\t"+gv)
+		}
+	} else {
+		for _, gv := range globals {
+			x = append(x, "var "+gv)
+		}
+	}
+
+	pf.Lines = append(pf.Lines[:end], append(x, pf.Lines[end:]...)...)
+	fileContent := strings.Join(pf.Lines, "\n")
+
+	return genny.NewFile(gf.Name(), strings.NewReader(fileContent)), nil
+}
+
+func findLastGlobal(f *ast.File, fset *token.FileSet, fileLines []string) (int, bool) {
+	var end = -1
+	var isGroup = false
+
+	ast.Inspect(f, func(n ast.Node) bool {
+		switch x := n.(type) {
+		case *ast.GenDecl:
+			if x.Tok == token.VAR && x.Rparen > 0 {
+				end = fset.Position(x.Rparen).Line - 1
+				isGroup = true
+				return false // if var block was found, break the loop
+			} else if x.Tok == token.IMPORT {
+				end = fset.Position(x.Rparen).Line
+				if end == 0 {
+					end = fset.Position(x.End()).Line
+				} // this is the last resort. just after the import block
+				return true
+			}
+
+		case *ast.ValueSpec:
+			end = fset.Position(x.End()).Line
+			return true // continue the block and find the last one
+
+		case *ast.FuncDecl:
+			return false // break the loop if the any function started
+		}
+		return true
+	})
+
+	return end, isGroup
+}

gogen/add_global_test.go

@@ -0,0 +1,165 @@
+package gogen
+
+import (
+	"io/ioutil"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/gobuffalo/genny/v2"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_AddGlobal(t *testing.T) {
+
+	tt := []struct {
+		Name     string
+		Source   string
+		Expected string
+	}{
+		{"single", globalBeforeSingle, globalAfterSingle},
+		{"block", globalBeforeBlock, globalAfterBlock},
+		{"empty1", globalBeforeEmpty1, globalAfterEmpty1},
+		{"empty2", globalBeforeEmpty2, globalAfterEmpty2},
+	}
+
+	for _, tc := range tt {
+		t.Run(tc.Name, func(st *testing.T) {
+			r := require.New(t)
+			path := filepath.Join("actions", "app.go")
+			f := genny.NewFile(path, strings.NewReader(tc.Source))
+
+			f, err := AddGlobal(f, "amount int", "top string")
+			r.NoError(err)
+
+			b, err := ioutil.ReadAll(f)
+			r.NoError(err)
+
+			r.Equal(tc.Expected, string(b))
+		})
+	}
+}
+
+const globalBeforeSingle = `package actions
+
+import (
+	"github.com/gobuffalo/buffalo"
+	"github.com/gobuffalo/envy"
+)
+
+var count int
+var chkme int
+func main() {
+	var local int
+	return
+}
+`
+
+const globalAfterSingle = `package actions
+
+import (
+	"github.com/gobuffalo/buffalo"
+	"github.com/gobuffalo/envy"
+)
+
+var count int
+var chkme int
+var amount int
+var top string
+func main() {
+	var local int
+	return
+}
+`
+
+const globalBeforeBlock = `package actions
+
+import (
+	"github.com/gobuffalo/buffalo"
+	"github.com/gobuffalo/envy"
+)
+
+var isOK bool
+
+var (
+	count int
+	chkme int
+
+)
+func main() {
+	var local int
+	return
+}
+`
+
+const globalAfterBlock = `package actions
+
+import (
+	"github.com/gobuffalo/buffalo"
+	"github.com/gobuffalo/envy"
+)
+
+var isOK bool
+
+var (
+	count int
+	chkme int
+
+	amount int
+	top string
+)
+func main() {
+	var local int
+	return
+}
+`
+
+const globalBeforeEmpty1 = `package actions
+
+import "github.com/gobuffalo/buffalo"
+import "github.com/gobuffalo/envy"
+
+func main() {
+	var local int
+	return
+}
+`
+
+const globalAfterEmpty1 = `package actions
+
+import "github.com/gobuffalo/buffalo"
+import "github.com/gobuffalo/envy"
+var amount int
+var top string
+
+func main() {
+	var local int
+	return
+}
+`
+
+const globalBeforeEmpty2 = `package actions
+
+import (
+	"github.com/gobuffalo/buffalo"
+	"github.com/gobuffalo/envy"
+)
+func main() {
+	var local int
+	return
+}
+`
+
+const globalAfterEmpty2 = `package actions
+
+import (
+	"github.com/gobuffalo/buffalo"
+	"github.com/gobuffalo/envy"
+)
+var amount int
+var top string
+func main() {
+	var local int
+	return
+}
+`

gogen/add_inside_block.go

@@ -54,24 +54,15 @@ func findBlockCoordinates(search string, pf ParsedFile) (int, int) {
 
 	ast.Inspect(pf.Ast, func(n ast.Node) bool {
 		switch x := n.(type) {
-		case *ast.StructType:
+		case *ast.StructType, *ast.BlockStmt:
 			line := pf.FileSet.Position(x.Pos()).Line
 			structDeclaration := fmt.Sprintf("%s\n", pf.Lines[line-1])
 
 			if strings.Contains(structDeclaration, search) {
 				start = line
-				end = pf.FileSet.Position(x.End()).Line
-				return false
+				end = pf.FileSet.Position(x.End()).Line - 1
+				return false // should return false to guarantee the result
 			}
-
-		case *ast.BlockStmt:
-			start = pf.FileSet.Position(x.Lbrace).Line
-			blockDeclaration := fmt.Sprintf("%s\n", pf.Lines[start-1])
-
-			if strings.Contains(blockDeclaration, search) {
-				end = pf.FileSet.Position(x.Rbrace).Line - 1
-			}
-
 		}
 		return true
 	})

gogen/replace_block.go

@@ -0,0 +1,37 @@
+package gogen
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/gobuffalo/genny/v2"
+)
+
+// ReplaceBlock replaces found block with given new open/close statements
+/*
+	open - start statement of the block e.g. `if isTrue {`
+	close - close statement of the block e.g. `}`
+	newOpen - new start statement of the block e.g. `defer func(isTrue) {`
+	newClose - new close statement of the block e.g. `}()`
+*/
+func ReplaceBlock(gf genny.File, open, close, newOpen, newClose string) (genny.File, error) {
+	pf, err := ParseFile(gf)
+	if err != nil {
+		return gf, err
+	}
+	gf = pf.File
+
+	start, end := findBlockCoordinates(open, pf)
+	if end < 0 {
+		return gf, fmt.Errorf("could not find desired block in %s", gf.Name())
+	}
+
+	lines := pf.Lines[:start-1]
+	lines = append(lines, strings.Replace(pf.Lines[start-1], open, newOpen, 1)) // open
+	lines = append(lines, pf.Lines[start:end]...)                               // body
+	lines = append(lines, strings.Replace(pf.Lines[end], close, newClose, 1))   // close
+	lines = append(lines, pf.Lines[end+1:]...)
+	fileContent := strings.Join(lines, "\n")
+
+	return genny.NewFile(gf.Name(), strings.NewReader(fileContent)), nil
+}

gogen/replace_block_body.go

@@ -7,7 +7,7 @@ import (
 	"github.com/gobuffalo/genny/v2"
 )
 
-// ReplaceBlockBody will replace found block with expressions passed
+// ReplaceBlockBody will replace found block body with expressions passed
 func ReplaceBlockBody(gf genny.File, search string, expressions ...string) (genny.File, error) {
 	pf, err := ParseFile(gf)
 	if err != nil {