From cd88a0edaa762602308538e482d46c78b6115655 Mon Sep 17 00:00:00 2001
From: nimodai <nimodai@tencent.com>
Date: Thu, 25 Mar 2021 16:31:01 +0800
Subject: [PATCH] fix: add include, and exclude is invalid if include is set

---
 cmd/swag/main.go |  6 ++++
 gen/gen.go       |  4 +++
 parser.go        | 82 ++++++++++++++++++++++++++++++++++++++----------
 3 files changed, 76 insertions(+), 16 deletions(-)

diff --git a/cmd/swag/main.go b/cmd/swag/main.go
index 75c175617..c290c53fc 100644
--- a/cmd/swag/main.go
+++ b/cmd/swag/main.go
@@ -13,6 +13,7 @@ import (
 const (
 	searchDirFlag        = "dir"
 	excludeFlag          = "exclude"
+	includeFlag          = "include"
 	generalInfoFlag      = "generalInfo"
 	propertyStrategyFlag = "propertyStrategy"
 	outputFlag           = "output"
@@ -38,6 +39,10 @@ var initFlags = []cli.Flag{
 		Value:   "./",
 		Usage:   "Directory you want to parse",
 	},
+	&cli.StringFlag{
+		Name:  includeFlag,
+		Usage: "include directories and files when searching, comma separated, exclude is invalid if include is set",
+	},
 	&cli.StringFlag{
 		Name:  excludeFlag,
 		Usage: "Exclude directories and files when searching, comma separated",
@@ -100,6 +105,7 @@ func initAction(c *cli.Context) error {
 
 	return gen.New().Build(&gen.Config{
 		SearchDir:           c.String(searchDirFlag),
+		Includes:            c.String(includeFlag),
 		Excludes:            c.String(excludeFlag),
 		MainAPIFile:         c.String(generalInfoFlag),
 		PropNamingStrategy:  strategy,
diff --git a/gen/gen.go b/gen/gen.go
index 291517485..89dd863ad 100644
--- a/gen/gen.go
+++ b/gen/gen.go
@@ -39,6 +39,9 @@ type Config struct {
 	// SearchDir the swag would be parse
 	SearchDir string
 
+	// includes dirs and files in SearchDir,comma separated
+	Includes string
+
 	// excludes dirs and files in SearchDir,comma separated
 	Excludes string
 
@@ -82,6 +85,7 @@ func (g *Gen) Build(config *Config) error {
 	log.Println("Generate swagger docs....")
 	p := swag.New(swag.SetMarkdownFileDirectory(config.MarkdownFilesDir),
 		swag.SetExcludedDirsAndFiles(config.Excludes),
+		swag.SetIncludedDirsAndFiles(config.Includes),
 		swag.SetCodeExamplesDirectory(config.CodeExampleFilesDir))
 	p.PropNamingStrategy = config.PropNamingStrategy
 	p.ParseVendor = config.ParseVendor
diff --git a/parser.go b/parser.go
index 4ca9bbfb3..c18a0d172 100644
--- a/parser.go
+++ b/parser.go
@@ -92,7 +92,10 @@ type Parser struct {
 	collectionFormatInQuery string
 
 	// excludes excludes dirs and files in SearchDir
-	excludes map[string]bool
+	excludes []string
+
+	// includes include dirs and files in SearchDir
+	includes []string
 }
 
 // New creates a new Parser with default properties.
@@ -120,7 +123,8 @@ func New(options ...func(*Parser)) *Parser {
 		outputSchemas:      make(map[*TypeSpecDef]*Schema),
 		existSchemaNames:   make(map[string]*Schema),
 		toBeRenamedSchemas: make(map[string]string),
-		excludes:           make(map[string]bool),
+		excludes:           make([]string, 0),
+		includes:           make([]string, 0),
 	}
 
 	for _, option := range options {
@@ -144,16 +148,29 @@ func SetCodeExamplesDirectory(directoryPath string) func(*Parser) {
 	}
 }
 
+func buildIncludeOrExclude(paths string) []string {
+	result := make([]string, 0)
+	for _, f := range strings.Split(paths, ",") {
+		f = strings.TrimSpace(f)
+		if f != "" {
+			f = filepath.Clean(f)
+			result = append(result, f)
+		}
+	}
+	return result
+}
+
 // SetExcludedDirsAndFiles sets directories and files to be excluded when searching
 func SetExcludedDirsAndFiles(excludes string) func(*Parser) {
 	return func(p *Parser) {
-		for _, f := range strings.Split(excludes, ",") {
-			f = strings.TrimSpace(f)
-			if f != "" {
-				f = filepath.Clean(f)
-				p.excludes[f] = true
-			}
-		}
+		p.excludes = buildIncludeOrExclude(excludes)
+	}
+}
+
+// SetIncludedDirsAndFiles sets directories and files to be include when searching
+func SetIncludedDirsAndFiles(includes string) func(*Parser) {
+	return func(p *Parser) {
+		p.includes = buildIncludeOrExclude(includes)
 	}
 }
 
@@ -557,6 +574,9 @@ func getSchemes(commentLine string) []string {
 
 // ParseRouterAPIInfo parses router api info for given astFile
 func (parser *Parser) ParseRouterAPIInfo(fileName string, astFile *ast.File) error {
+	if !parser.checkFileNeedParse(fileName) {
+		return nil
+	}
 	for _, astDescription := range astFile.Decls {
 		switch astDeclaration := astDescription.(type) {
 		case *ast.FuncDecl:
@@ -598,6 +618,43 @@ func (parser *Parser) ParseRouterAPIInfo(fileName string, astFile *ast.File) err
 	return nil
 }
 
+//checkFileNeedParse Check if this file needs to be parse depend on include and exclude.
+// return true is this file need to be parse
+func (parser *Parser) checkFileNeedParse(fileName string) bool {
+	if len(parser.includes) > 0 {
+		if parser.checkFileInIncludes(fileName) {
+			return true
+		}
+		return false
+	}
+	if parser.checkFileInExcludes(fileName) {
+		return false
+	}
+	return true
+}
+
+// checkFileInIncludes check this fileName in includes slice
+// return true if fileName is prefix of one of includes path
+func (parser *Parser) checkFileInIncludes(fileName string) bool {
+	for _, include := range parser.includes {
+		if strings.HasPrefix(fileName, include) {
+			return true
+		}
+	}
+	return false
+}
+
+// checkFileInExcludes check this fileName in excludes slice
+// return true if fileName is prefix of one of excludes path
+func (parser *Parser) checkFileInExcludes(fileName string) bool {
+	for _, exclude := range parser.excludes {
+		if strings.HasPrefix(fileName, exclude) {
+			return true
+		}
+	}
+	return false
+}
+
 func convertFromSpecificToPrimitive(typeName string) (string, error) {
 	name := typeName
 	if strings.ContainsRune(name, '.') {
@@ -1434,14 +1491,7 @@ func (parser *Parser) Skip(path string, f os.FileInfo) error {
 			len(f.Name()) > 1 && f.Name()[0] == '.' { // exclude all hidden folder
 			return filepath.SkipDir
 		}
-
-		if parser.excludes != nil {
-			if _, ok := parser.excludes[path]; ok {
-				return filepath.SkipDir
-			}
-		}
 	}
-
 	return nil
 }