support -typed option (#92)

sanposhiho · Sep 4, 2024 · b142dd4 · b142dd4
1 parent 5e6d32a
commit b142dd4
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 4 deletions.
diff --git a/gomockhandler.go b/gomockhandler.go
@@ -40,6 +40,7 @@ var (
 	buildFlags      = flag.String("build_flags", "", "[option for configure mockgen] (reflect mode) Additional flags for go build.")
 	progOnly        = flag.Bool("prog_only", false, "[option for configure mockgen] (reflect mode) Only generate the reflection program; write it to stdout and exit.")
 	debugParser     = flag.Bool("debug_parser", false, "[option for configure mockgen] Print out parser results only.")
+	typed           = flag.Bool("typed", false, "[option for configure mockgen] Generate Type-safe 'Return', 'Do', 'DoAndReturn' function.")
 )
 
 func main() {
@@ -106,11 +107,11 @@ func prepareMockgenRunner() mockgen.Runner {
 				log.Fatalf("Parse package name failed: %v", err)
 			}
 		}
-		return reflectmode.NewRunner(packageName, interfaces, *source, *destination, *packageOut, *imports, *auxFiles, *buildFlags, *mockNames, *selfPackage, *copyrightFile, *execOnly, *progOnly, *writePkgComment, *debugParser)
+		return reflectmode.NewRunner(packageName, interfaces, *source, *destination, *packageOut, *imports, *auxFiles, *buildFlags, *mockNames, *selfPackage, *copyrightFile, *execOnly, *progOnly, *writePkgComment, *debugParser, *typed)
 	}
 
 	// source mode
-	return sourcemode.NewRunner(*source, *destination, *packageOut, *imports, *auxFiles, *mockNames, *selfPackage, *copyrightFile, *writePkgComment, *debugParser)
+	return sourcemode.NewRunner(*source, *destination, *packageOut, *imports, *auxFiles, *mockNames, *selfPackage, *copyrightFile, *writePkgComment, *debugParser, *typed)
 }
 
 // Plundered from golang/mock/mockgen/parse.go.

diff --git a/internal/mockgen/reflectmode/runner.go b/internal/mockgen/reflectmode/runner.go
@@ -21,9 +21,10 @@ type Runner struct {
 	ProgOnly        *bool  `json:"prog_only,omitempty"`
 	WritePkgComment *bool  `json:"write_pkg_comment,omitempty"`
 	DebugParser     *bool  `json:"debug_parser,omitempty"`
+	Typed           *bool  `json:"typed,omitempty"`
 }
 
-func NewRunner(pn, ifs, source, dest, pkg, imp, af, bf, mn, spkg, cf, eo string, po, wpc, dp bool) *Runner {
+func NewRunner(pn, ifs, source, dest, pkg, imp, af, bf, mn, spkg, cf, eo string, po, wpc, dp, typed bool) *Runner {
 	var wpcp *bool
 	if wpc != true {
 		// The default value of wpc is true
@@ -42,6 +43,12 @@ func NewRunner(pn, ifs, source, dest, pkg, imp, af, bf, mn, spkg, cf, eo string,
 		pop = &po
 	}
 
+	var tp *bool
+	if typed != false {
+		// The default value of typed is false
+		tp = &typed
+	}
+
 	return &Runner{
 		PackageName:     pn,
 		Interfaces:      ifs,
@@ -58,6 +65,7 @@ func NewRunner(pn, ifs, source, dest, pkg, imp, af, bf, mn, spkg, cf, eo string,
 		ProgOnly:        pop,
 		WritePkgComment: wpcp,
 		DebugParser:     dpp,
+		Typed:           tp,
 	}
 }
 
@@ -128,5 +136,8 @@ func (r *Runner) options() []string {
 	if r.ProgOnly != nil {
 		opts = append(opts, "-prog_only="+strconv.FormatBool(*r.ProgOnly))
 	}
+	if r.Typed != nil {
+		opts = append(opts, "-typed="+strconv.FormatBool(*r.Typed))
+	}
 	return opts
 }
diff --git a/internal/mockgen/sourcemode/runner.go b/internal/mockgen/sourcemode/runner.go
@@ -16,9 +16,10 @@ type Runner struct {
 	CopyrightFile   string `json:"copyright_file,omitempty"`
 	WritePkgComment *bool  `json:"write_pkg_comment,omitempty"`
 	DebugParser     *bool  `json:"debug_parser,omitempty"`
+	Typed           *bool  `json:"typed,omitempty"`
 }
 
-func NewRunner(source, dest, pkg, imp, af, mn, spkg, cf string, wpc, dp bool) *Runner {
+func NewRunner(source, dest, pkg, imp, af, mn, spkg, cf string, wpc, dp, typed bool) *Runner {
 	var wpcp *bool
 	if wpc != true {
 		// The default value of wpc is true
@@ -30,6 +31,13 @@ func NewRunner(source, dest, pkg, imp, af, mn, spkg, cf string, wpc, dp bool) *R
 		// The default value of dp is false
 		dpp = &dp
 	}
+
+	var tp *bool
+	if typed != false {
+		// The default value of typed is false
+		tp = &typed
+	}
+
 	return &Runner{
 		Source:          source,
 		Destination:     dest,
@@ -41,6 +49,7 @@ func NewRunner(source, dest, pkg, imp, af, mn, spkg, cf string, wpc, dp bool) *R
 		CopyrightFile:   cf,
 		WritePkgComment: wpcp,
 		DebugParser:     dpp,
+		Typed:           tp,
 	}
 }
 
@@ -100,6 +109,9 @@ func (r *Runner) options() []string {
 	if r.DebugParser != nil {
 		opts = append(opts, "-debug_parser="+strconv.FormatBool(*r.DebugParser))
 	}
+	if r.Typed != nil {
+		opts = append(opts, "-typed="+strconv.FormatBool(*r.Typed))
+	}
 
 	return opts
 }