separate go:generate logic from counterfeiter:generate logic

command/runner.go

@@ -6,17 +6,23 @@ import (
 	"io/ioutil"
 	"os"
 	"path/filepath"
-	"regexp"
 	"sort"
 	"strconv"
 	"strings"
 )
 
 func Detect(cwd string, args []string, generateMode bool) ([]Invocation, error) {
-	if generateMode || invokedByGoGenerate() {
-		return invocations(cwd, generateMode)
+	if generateMode {
+		return generateModeInvocations(cwd)
 	}
-	i, err := NewInvocation("", 0, args)
+
+	file := os.Getenv("GOFILE")
+	var lineno int
+	if goline, err := strconv.Atoi(os.Getenv("GOLINE")); err == nil {
+		lineno = goline
+	}
+
+	i, err := NewInvocation(file, lineno, args)
 	if err != nil {
 		return nil, err
 	}
@@ -41,11 +47,7 @@ func NewInvocation(file string, line int, args []string) (Invocation, error) {
 	return i, nil
 }
 
-func invokedByGoGenerate() bool {
-	return os.Getenv("DOLLAR") == "$"
-}
-
-func invocations(cwd string, generateMode bool) ([]Invocation, error) {
+func generateModeInvocations(cwd string) ([]Invocation, error) {
 	var result []Invocation
 	// Find all the go files
 	pkg, err := build.ImportDir(cwd, build.IgnoreVendor)
@@ -59,65 +61,19 @@ func invocations(cwd string, generateMode bool) ([]Invocation, error) {
 	gofiles = append(gofiles, pkg.TestGoFiles...)
 	gofiles = append(gofiles, pkg.XTestGoFiles...)
 	sort.Strings(gofiles)
-	var line int
-	if !generateMode {
-		// generateMode means counterfeiter:generate, not go:generate
-		line, err = strconv.Atoi(os.Getenv("GOLINE"))
-		if err != nil {
-			return nil, err
-		}
-	}
 
-	for i := range gofiles {
-		i, err := open(cwd, gofiles[i], generateMode)
+	for _, file := range gofiles {
+		invocations, err := invocationsInFile(cwd, file)
 		if err != nil {
 			return nil, err
 		}
-		result = append(result, i...)
-		if generateMode {
-			continue
-		}
-		if len(result) > 0 && result[0].File != os.Getenv("GOFILE") {
-			return nil, nil
-		}
-		if len(result) > 0 && result[0].Line != line {
-			return nil, nil
-		}
+		result = append(result, invocations...)
 	}
 
 	return result, nil
 }
 
-var directive = regexp.MustCompile(`(?mi)^//(go:generate|counterfeiter:generate)\s*(.*)?\s*$`)
-var args = regexp.MustCompile(`(?mi)^(?:go run github\.com/maxbrunsfeld/counterfeiter/v6|gobin -m -run github\.com/maxbrunsfeld/counterfeiter/v6|counterfeiter|counterfeiter.exe)\s*(.*)?\s*$`)
-
-type match struct {
-	directive string
-	args      []string
-}
-
-func matchForString(s string) *match {
-	m := directive.FindStringSubmatch(s)
-	if m == nil {
-		return nil
-	}
-	if m[1] == "counterfeiter:generate" {
-		return &match{
-			directive: m[1],
-			args:      stringToArgs(m[2]),
-		}
-	}
-	m2 := args.FindStringSubmatch(m[2])
-	if m2 == nil {
-		return nil
-	}
-	return &match{
-		directive: m[1],
-		args:      stringToArgs(m2[1]),
-	}
-}
-
-func open(dir string, file string, generateMode bool) ([]Invocation, error) {
+func invocationsInFile(dir string, file string) ([]Invocation, error) {
 	str, err := ioutil.ReadFile(filepath.Join(dir, file))
 	if err != nil {
 		return nil, err
@@ -128,30 +84,30 @@ func open(dir string, file string, generateMode bool) ([]Invocation, error) {
 	line := 0
 	for i := range lines {
 		line++
-		match := matchForString(lines[i])
-		if match == nil {
+		args, ok := matchForString(lines[i])
+		if !ok {
 			continue
 		}
-		inv, err := NewInvocation(file, line, match.args)
+		inv, err := NewInvocation(file, line, args)
 		if err != nil {
 			return nil, err
 		}
 
-		if generateMode && match.directive == "counterfeiter:generate" {
-			result = append(result, inv)
-		}
-
-		if !generateMode && match.directive == "go:generate" {
-			if len(inv.Args) == 2 && strings.EqualFold(strings.TrimSpace(inv.Args[1]), "-generate") {
-				continue
-			}
-			result = append(result, inv)
-		}
+		result = append(result, inv)
 	}
 
 	return result, nil
 }
 
+const generateDirectivePrefix = "//counterfeiter:generate "
+
+func matchForString(s string) ([]string, bool) {
+	if !strings.HasPrefix(s, generateDirectivePrefix) {
+		return nil, false
+	}
+	return stringToArgs(s[len(generateDirectivePrefix):]), true
+}
+
 func stringToArgs(s string) []string {
 	a := strings.Fields(s)
 	result := []string{

command/runner_internals_test.go

@@ -28,13 +28,11 @@ func testRegexp(t *testing.T, when spec.G, it spec.S) {
 		cases = []Case{
 			{
 				input:   "//go:generate counterfeiter . Intf",
-				matches: true,
-				args:    []string{".", "Intf"},
+				matches: false,
 			},
 			{
 				input:   "//go:generate go run github.com/maxbrunsfeld/counterfeiter/v6 . Intf",
-				matches: true,
-				args:    []string{".", "Intf"},
+				matches: false,
 			},
 			{
 				input:   "//counterfeiter:generate . Intf",
@@ -44,7 +42,6 @@ func testRegexp(t *testing.T, when spec.G, it spec.S) {
 			{
 				input:   "//go:generate  stringer -type=Enum",
 				matches: false,
-				args:    []string{".", "Intf"},
 			},
 		}
 	})
@@ -56,12 +53,12 @@ func testRegexp(t *testing.T, when spec.G, it spec.S) {
 
 	it("matches lines appropriately", func() {
 		for _, c := range cases {
-			result := matchForString(c.input)
+			result, ok := matchForString(c.input)
 			if c.matches {
-				Expect(result).NotTo(BeNil(), c.input)
-				Expect(result.args).To(ConsistOf(c.args))
+				Expect(ok).To(BeTrue())
+				Expect(result).To(ConsistOf(c.args))
 			} else {
-				Expect(result).To(BeNil(), c.input)
+				Expect(ok).To(BeFalse())
 			}
 		}
 	})

command/runner_test.go

@@ -91,29 +91,5 @@ func testRunner(t *testing.T, when spec.G, it spec.S) {
 			Expect(i[0].Args[1]).To(Equal("."))
 			Expect(i[0].Args[2]).To(Equal("AliasedInterface"))
 		})
-
-		when("there is a mismatch in the file name", func() {
-			it.Before(func() {
-				os.Setenv("GOFILE", "some_other_file.go")
-			})
-
-			it("has no invocations", func() {
-				i, err := command.Detect(filepath.Join(".", "..", "fixtures"), []string{"counterfeiter", ".", "AliasedInterface"}, false)
-				Expect(err).NotTo(HaveOccurred())
-				Expect(i).To(HaveLen(0))
-			})
-		})
-
-		when("there is a mismatch in the line number", func() {
-			it.Before(func() {
-				os.Setenv("GOLINE", "100")
-			})
-
-			it("has no invocations", func() {
-				i, err := command.Detect(filepath.Join(".", "..", "fixtures"), []string{"counterfeiter", ".", "AliasedInterface"}, false)
-				Expect(err).NotTo(HaveOccurred())
-				Expect(i).To(HaveLen(0))
-			})
-		})
 	})
 }