Merge branch 'main' into dev/keyi_string_json

tool/trimmer/main.go

@@ -20,13 +20,12 @@ import (
 	"path/filepath"
 	"strings"
 
+	"github.com/cloudwego/thriftgo/generator"
 	"github.com/cloudwego/thriftgo/parser"
 	"github.com/cloudwego/thriftgo/semantic"
 	"github.com/cloudwego/thriftgo/tool/trimmer/dump"
 	"github.com/cloudwego/thriftgo/tool/trimmer/trim"
 	"github.com/cloudwego/thriftgo/version"
-
-	"github.com/cloudwego/thriftgo/generator"
 )
 
 var (
@@ -99,7 +98,7 @@ func main() {
 					os.Exit(2)
 				}
 			} else {
-				println("-o should be set as a valid dir to enable -r", err.Error())
+				println("-o should be set as a valid dir to enable -r")
 				os.Exit(2)
 			}
 		}

tool/trimmer/test_cases/sample1.thrift

@@ -53,6 +53,7 @@ enum Gender {
     FEMALE (key = "1", key = "2", key2 = "v2")
 } (a = "b")
 
+#@PRESERvE
 struct Address {
     1: required string(key = "v") street
     2: required string city

tool/trimmer/test_cases/tests/example2.thrift

@@ -22,6 +22,6 @@ struct E{
 
 }
 
-service MyService{
+service MyService extends sample1.EmployeeService{
     string A(1:required E req,2:required test.TestStruct t)
 }

tool/trimmer/trim/mark.go

@@ -49,14 +49,17 @@ func (t *Trimmer) markService(svc *parser.Service, ast *parser.Thrift, filename
 					t.marks[filename][svc] = true
 					t.markFunction(function, ast, filename)
 					t.trimMethodValid[i] = true
-					continue
 				}
 			}
 			continue
 		}
 		t.markFunction(function, ast, filename)
 	}
 
+	if t.trimMethods != nil && (svc.Extends != "" || svc.Reference != nil) {
+		t.traceExtendMethod(svc, svc, ast, filename)
+	}
+
 	if svc.Extends != "" && t.marks[filename][svc] {
 		// handle extension
 		if svc.Reference != nil {
@@ -169,3 +172,48 @@ func (t *Trimmer) markTypeDef(theType *parser.Type, ast *parser.Thrift, filename
 		}
 	}
 }
+
+// for -m, trace the extends and find specified method to base on
+func (t *Trimmer) traceExtendMethod(father, svc *parser.Service, ast *parser.Thrift, filename string) (ret bool) {
+	for _, function := range svc.Functions {
+		funcName := father.Name + "." + function.Name
+		for i, method := range t.trimMethods {
+			if funcName == method {
+				t.marks[filename][svc] = true
+				t.markFunction(function, ast, filename)
+				t.trimMethodValid[i] = true
+				ret = true
+			}
+		}
+	}
+	if svc.Extends != "" {
+		var nextSvc *parser.Service
+		var nextAst *parser.Thrift
+		if svc.Reference == nil {
+			for i, extend := range ast.Services {
+				if extend.Name == svc.Extends {
+					nextSvc = ast.Services[i]
+					nextAst = ast
+					break
+				}
+			}
+		} else {
+			for i, extend := range ast.Includes[svc.Reference.Index].Reference.Services {
+				if extend.Name == svc.Reference.Name {
+					nextSvc = ast.Includes[svc.Reference.Index].Reference.Services[i]
+					nextAst = ast.Includes[svc.Reference.Index].Reference
+					break
+				}
+			}
+		}
+		back := t.traceExtendMethod(father, nextSvc, nextAst, filename)
+		ret = back || ret
+	}
+	if ret {
+		t.marks[filename][svc] = true
+		if svc.Reference != nil {
+			t.marks[filename][ast.Includes[svc.Reference.Index]] = true
+		}
+	}
+	return ret
+}

tool/trimmer/trim/traversal.go

@@ -14,7 +14,12 @@
 
 package trim
 
-import "github.com/cloudwego/thriftgo/parser"
+import (
+	"regexp"
+	"strings"
+
+	"github.com/cloudwego/thriftgo/parser"
+)
 
 // traverse and remove the unmarked part of ast
 func (t *Trimmer) traversal(ast *parser.Thrift, filename string) {
@@ -30,7 +35,7 @@ func (t *Trimmer) traversal(ast *parser.Thrift, filename string) {
 
 	var listStruct []*parser.StructLike
 	for i := range ast.Structs {
-		if t.marks[filename][ast.Structs[i]] {
+		if t.marks[filename][ast.Structs[i]] || checkPreserve(ast.Structs[i]) {
 			listStruct = append(listStruct, ast.Structs[i])
 		}
 	}
@@ -69,3 +74,9 @@ func (t *Trimmer) traversal(ast *parser.Thrift, filename string) {
 	}
 	ast.Services = listService
 }
+
+func checkPreserve(theStruct *parser.StructLike) bool {
+	pattern := `^[\s]*(\/\/|#)[\s]*@preserve[\s]*$`
+	regex := regexp.MustCompile(pattern)
+	return regex.MatchString(strings.ToLower(theStruct.ReservedComments))
+}