From 347186b68dc5961353de861b0551af9fba16fad6 Mon Sep 17 00:00:00 2001
From: "yuxuan.wang1" <yuxuan.wang1@bytedance.com>
Date: Thu, 9 Jan 2025 09:35:02 +0800
Subject: [PATCH] online: support scan_dir for custom template

---
 .../internal_pkg/generator/custom_template.go | 56 +++++++++++++++++++
 tool/internal_pkg/generator/generator_test.go | 29 ++++++++++
 .../internal_pkg/generator/test_dir/test0.txt |  0
 .../generator/test_dir/test_dir1/test1.txt    |  0
 4 files changed, 85 insertions(+)
 create mode 100644 tool/internal_pkg/generator/test_dir/test0.txt
 create mode 100644 tool/internal_pkg/generator/test_dir/test_dir1/test1.txt

diff --git a/tool/internal_pkg/generator/custom_template.go b/tool/internal_pkg/generator/custom_template.go
index 3f364589db..7d2aabd1fb 100644
--- a/tool/internal_pkg/generator/custom_template.go
+++ b/tool/internal_pkg/generator/custom_template.go
@@ -16,6 +16,7 @@ package generator
 
 import (
 	"fmt"
+	"io/fs"
 	"os"
 	"path"
 	"path/filepath"
@@ -47,6 +48,8 @@ type Update struct {
 	AppendTpl string `yaml:"append_tpl,omitempty"`
 	// Append import template. Use it to render import content to append.
 	ImportTpl []string `yaml:"import_tpl,omitempty"`
+	// Using this field to scan the current directory and all the subdirecotries recursively.
+	ScanDirRecursively bool `yaml:"scan_dir_recursively,omitempty"`
 }
 
 type Template struct {
@@ -97,6 +100,18 @@ func (c *customGenerator) loopGenerate(tpl *Template) error {
 		if util.Exists(filePath) && updateType(tpl.UpdateBehavior.Type) == skip {
 			continue
 		}
+		// using scanDir to substitute the default directory
+		scanPath, scanErr := scanDir(filePath, tpl.UpdateBehavior)
+		if scanErr != nil {
+			return scanErr
+		}
+		if scanPath != "" {
+			if updateType(tpl.UpdateBehavior.Type) == skip {
+				continue
+			}
+			filePath = scanPath
+		}
+
 		task := &Task{
 			Name: path.Base(renderPath),
 			Path: filePath,
@@ -111,6 +126,36 @@ func (c *customGenerator) loopGenerate(tpl *Template) error {
 	return nil
 }
 
+// scanDir scans the directory of the filePath recursively and returns the path of the file with the same name.
+func scanDir(filePath string, behavior *Update) (res string, err error) {
+	if behavior == nil || !behavior.ScanDirRecursively {
+		return "", nil
+	}
+	fileName := filepath.Base(filePath)
+	dirPath := filepath.Dir(filePath)
+	var found bool
+	if err := filepath.Walk(dirPath, func(path string, info fs.FileInfo, walkErr error) error {
+		if walkErr != nil {
+			return walkErr
+		}
+		if !info.IsDir() && info.Name() == fileName {
+			found = true
+			res = path
+			return nil
+		}
+		return nil
+	}); err != nil {
+		if os.IsNotExist(err) {
+			return "", nil
+		}
+		return "", err
+	}
+	if found {
+		return res, nil
+	}
+	return "", nil
+}
+
 func (c *customGenerator) commonGenerate(tpl *Template) error {
 	// Use all services including base service.
 	tmp := c.pkg.Methods
@@ -132,6 +177,17 @@ func (c *customGenerator) commonGenerate(tpl *Template) error {
 		log.Infof("skip generate file %s", tpl.Path)
 		return nil
 	}
+	// using scanDir to substitute the default directory
+	scanPath, scanErr := scanDir(filePath, tpl.UpdateBehavior)
+	if scanErr != nil {
+		return scanErr
+	}
+	if scanPath != "" {
+		if updateType(tpl.UpdateBehavior.Type) == skip {
+			return nil
+		}
+		filePath = scanPath
+	}
 	var f *File
 	if update && updateType(tpl.UpdateBehavior.Type) == incrementalUpdate {
 		cc := &commonCompleter{
diff --git a/tool/internal_pkg/generator/generator_test.go b/tool/internal_pkg/generator/generator_test.go
index 7b3d160a70..cff90d6ae3 100644
--- a/tool/internal_pkg/generator/generator_test.go
+++ b/tool/internal_pkg/generator/generator_test.go
@@ -231,3 +231,32 @@ func Test_needCallOpt(t *testing.T) {
 		test.Assert(t, needCallOpt(pkg))
 	})
 }
+
+func Test_scanDir(t *testing.T) {
+	testcases := []struct {
+		filePath   string
+		update     *Update
+		expectPath string
+		expectErr  error
+	}{
+		{
+			filePath:   "./test_dir/test1.txt",
+			update:     &Update{ScanDirRecursively: true, Type: "skip"},
+			expectPath: "test_dir/test_dir1/test1.txt",
+		},
+		{
+			filePath: "./test_dir/test2.txt",
+			update:   &Update{ScanDirRecursively: true, Type: "skip"},
+		},
+		{
+			filePath: "./non_exist_dir/test3.txt",
+			update:   &Update{ScanDirRecursively: true, Type: "skip"},
+		},
+	}
+
+	for _, tc := range testcases {
+		res, err := scanDir(tc.filePath, tc.update)
+		test.Assert(t, err == tc.expectErr, err)
+		test.Assert(t, res == tc.expectPath, res)
+	}
+}
diff --git a/tool/internal_pkg/generator/test_dir/test0.txt b/tool/internal_pkg/generator/test_dir/test0.txt
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tool/internal_pkg/generator/test_dir/test_dir1/test1.txt b/tool/internal_pkg/generator/test_dir/test_dir1/test1.txt
new file mode 100644
index 0000000000..e69de29bb2