feat: add yaml-config support for trimmer (#131)

Co-authored-by: Li2CO3 <[email protected]>

generator/golang/backend.go

@@ -86,7 +86,7 @@ func (g *GoBackend) Generate(req *plugin.Request, log backend.LogFunc) *plugin.R
 	g.log = log
 	g.prepareUtilities()
 	if g.utils.Features().TrimIDL {
-		err := trim.TrimAST(req.AST, nil, false)
+		err := trim.TrimAST(&trim.TrimASTArg{Ast: req.AST, TrimMethods: nil, Preserve: nil})
 		if err != nil {
 			g.log.Warn("trim error:", err.Error())
 		}

tool/trimmer/args.go

@@ -63,8 +63,8 @@ func (a *Arguments) BuildFlags() *flag.FlagSet {
 	f.Var(&a.Methods, "m", "")
 	f.Var(&a.Methods, "method", "")
 
-	f.StringVar(&a.Preserve, "p", "true", "")
-	f.StringVar(&a.Preserve, "preserve", "true", "")
+	f.StringVar(&a.Preserve, "p", "", "")
+	f.StringVar(&a.Preserve, "preserve", "", "")
 
 	f.Usage = help
 	return f

tool/trimmer/main.go

@@ -53,13 +53,14 @@ func main() {
 		os.Exit(0)
 	}
 
-	preserve := true
+	var preserveInput *bool
 	if a.Preserve != "" {
-		preserve, err = strconv.ParseBool(a.Preserve)
+		preserve, err := strconv.ParseBool(a.Preserve)
 		if err != nil {
 			help()
 			os.Exit(2)
 		}
+		preserveInput = &preserve
 	}
 
 	// parse file to ast
@@ -74,7 +75,9 @@ func main() {
 	check(semantic.ResolveSymbols(ast))
 
 	// trim ast
-	check(trim.TrimAST(ast, a.Methods, !preserve))
+	check(trim.TrimAST(&trim.TrimASTArg{
+		Ast: ast, TrimMethods: a.Methods, Preserve: preserveInput,
+	}))
 
 	// dump the trimmed ast to idl
 	idl, err := dump.DumpIDL(ast)
@@ -117,16 +120,7 @@ func main() {
 			println("output-dir should be set outside of -r base-dir to avoid overlay")
 			os.Exit(2)
 		}
-		createDirTree(a.Recurse, a.OutputFile)
 		recurseDump(ast, a.Recurse, a.OutputFile)
-		relativePath, err := filepath.Rel(a.Recurse, a.IDL)
-		if err != nil {
-			println("-r input err, range should cover all the target IDLs;", err.Error())
-			os.Exit(2)
-		}
-		outputFileUrl := filepath.Join(a.OutputFile, relativePath)
-		check(writeStringToFile(outputFileUrl, idl))
-		removeEmptyDir(a.OutputFile)
 	} else {
 		check(writeStringToFile(a.OutputFile, idl))
 	}
@@ -139,16 +133,21 @@ func recurseDump(ast *parser.Thrift, sourceDir, outDir string) {
 	if ast == nil {
 		return
 	}
+	out, err := dump.DumpIDL(ast)
+	check(err)
+	relativeUrl, err := filepath.Rel(sourceDir, ast.Filename)
+	if err != nil {
+		println("-r input err, range should cover all the target IDLs;", err.Error())
+		os.Exit(2)
+	}
+	outputFileUrl := filepath.Join(outDir, relativeUrl)
+	err = os.MkdirAll(filepath.Dir(outputFileUrl), os.ModePerm)
+	if err != nil {
+		println("mkdir", filepath.Dir(outputFileUrl), "error:", err.Error())
+		os.Exit(2)
+	}
+	check(writeStringToFile(outputFileUrl, out))
 	for _, includes := range ast.Includes {
-		out, err := dump.DumpIDL(includes.Reference)
-		check(err)
-		relativeUrl, err := filepath.Rel(sourceDir, includes.Reference.Filename)
-		if err != nil {
-			println("-r input err, range should cover all the target IDLs;", err.Error())
-			os.Exit(2)
-		}
-		outputFileUrl := filepath.Join(outDir, relativeUrl)
-		check(writeStringToFile(outputFileUrl, out))
 		recurseDump(includes.Reference, sourceDir, outDir)
 	}
 }

tool/trimmer/test_cases/tests/dir/dir2/test.thrift

@@ -13,8 +13,19 @@
 // limitations under the License.
 
 include "../../../sample1.thrift"
+include "../dir3/dir4/another.thrift"
 
 // @preserve
 struct TestStruct{
     1: sample1.Person person
+    2: another.AnotherStruct another
+}
+
+service TestService{
+    void func1()
+    another.AnotherStruct func2()
+    void func3()
+}
+
+union useless{
 }

tool/trimmer/test_cases/tests/dir/dir2/trim_config.yaml

@@ -0,0 +1,19 @@
+# Copyright 2023 CloudWeGo Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+methods:
+  - "TestService.func1"
+  - "TestService.func3"
+preserve: true
+preserved_structs:
+  - "useless"

tool/trimmer/test_cases/tests/dir/dir3/dir4/another.thrift

@@ -0,0 +1,16 @@
+// Copyright 2023 CloudWeGo Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+struct AnotherStruct{
+}

tool/trimmer/trim/config.go

@@ -0,0 +1,50 @@
+// Copyright 2023 CloudWeGo Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package trim
+
+import (
+	"fmt"
+	"os"
+	"path/filepath"
+
+	"gopkg.in/yaml.v3"
+)
+
+var DefaultYamlFileName = "trim_config.yaml"
+
+type YamlArguments struct {
+	Methods          []string `yaml:"methods,omitempty"`
+	Preserve         *bool    `yaml:"preserve,omitempty"`
+	PreservedStructs []string `yaml:"preserved_structs,omitempty"`
+}
+
+func ParseYamlConfig(path string) *YamlArguments {
+	cfg := YamlArguments{}
+	dataBytes, err := os.ReadFile(filepath.Join(path, DefaultYamlFileName))
+	if err != nil {
+		return nil
+	}
+	fmt.Println("using trim config:", filepath.Join(path, DefaultYamlFileName))
+	err = yaml.Unmarshal(dataBytes, &cfg)
+	if err != nil {
+		fmt.Println("unmarshal yaml config fail:", err)
+		return nil
+	}
+	if cfg.Preserve == nil {
+		t := true
+		cfg.Preserve = &t
+	}
+	return &cfg
+}

tool/trimmer/trim/mark.go

@@ -261,5 +261,10 @@ func (t *Trimmer) checkPreserve(theStruct *parser.StructLike) bool {
 	if t.forceTrimming {
 		return false
 	}
+	for _, name := range t.preservedStructs {
+		if name == theStruct.Name {
+			return true
+		}
+	}
 	return t.preserveRegex.MatchString(strings.ToLower(theStruct.ReservedComments))
 }

tool/trimmer/trim/traversal.go

@@ -40,15 +40,15 @@ func (t *Trimmer) traversal(ast *parser.Thrift, filename string) {
 
 	var listUnion []*parser.StructLike
 	for i := range ast.Unions {
-		if t.marks[filename][ast.Unions[i]] {
+		if t.marks[filename][ast.Unions[i]] || t.checkPreserve(ast.Unions[i]) {
 			listUnion = append(listUnion, ast.Unions[i])
 		}
 	}
 	ast.Unions = listUnion
 
 	var listException []*parser.StructLike
 	for i := range ast.Exceptions {
-		if t.marks[filename][ast.Exceptions[i]] {
+		if t.marks[filename][ast.Exceptions[i]] || t.checkPreserve(ast.Exceptions[i]) {
 			listException = append(listException, ast.Exceptions[i])
 		}
 	}