feat(tool): add env for using prutal to marshal (#1759)

Author: xiaost

tool/internal_pkg/generator/generator.go

@@ -605,11 +605,11 @@ func (g *generator) setImports(name string, pkg *PackageInfo) {
 				pkg.AddImports("fmt")
 			}
 			if m.GenArgResultStruct {
-				if env.UseProtoc() {
-					pkg.AddImports("proto")
-				} else {
+				if env.UsePrutalMarshal() {
 					// reuse "proto" so that no need to change code templates ...
 					pkg.AddImport("proto", "github.com/cloudwego/prutal")
+				} else {
+					pkg.AddImports("proto")
 				}
 			} else {
 				// for method Arg and Result

tool/internal_pkg/util/env/env.go

@@ -27,7 +27,19 @@ import (
 func UseProtoc() bool {
 	v := os.Getenv("KITEX_TOOL_USE_PROTOC")
 	if v == "" {
-		return false // disable protoc by default
+		return false // false by default
+	}
+	ok, _ := strconv.ParseBool(v)
+	return ok
+}
+
+func UsePrutalMarshal() bool {
+	if !UseProtoc() {
+		return true // if not using protoc, can only use prutal
+	}
+	v := os.Getenv("KITEX_TOOL_USE_PRUTAL_MARSHAL")
+	if v == "" {
+		return false // false by default
 	}
 	ok, _ := strconv.ParseBool(v)
 	return ok

tool/internal_pkg/util/env/env_test.go

@@ -37,3 +37,19 @@ func TestUseProtoc(t *testing.T) {
 	use = UseProtoc()
 	test.Assert(t, use == false, use)
 }
+
+func TestUsePrutalMarshal(t *testing.T) {
+	use := UsePrutalMarshal() // default: true
+	test.Assert(t, use == true, use)
+
+	t.Setenv("KITEX_TOOL_USE_PROTOC", "0")
+	use = UsePrutalMarshal()
+	test.Assert(t, use == true, use)
+
+	t.Setenv("KITEX_TOOL_USE_PROTOC", "1")
+	use = UsePrutalMarshal()
+	test.Assert(t, use == false, use)
+	t.Setenv("KITEX_TOOL_USE_PRUTAL_MARSHAL", "1")
+	use = UsePrutalMarshal()
+	test.Assert(t, use == true, use)
+}