Merge pull request #11 from snamiki1212/fix/include-flag

feat: add include flag

README.md

@@ -60,6 +60,7 @@ Flags:
   -e, --entity string     target entity name. e.g. --entity=User or --entity=*User
   -x, --exclude strings   exclude lo method e.g. --exclude=Map,Filter
   -h, --help              help for go-gen-lo
+  -n, --include strings   include lo method with regex e.g. --include=Filter*,Map
   -i, --input string      input file name
   -o, --output string     output file name
   -r, --rename strings    rename method e.g. --rename=Map:Loop

cmd/root.go

@@ -94,6 +94,9 @@ func init() {
 	// exclude
 	rootCmd.Flags().StringSliceVarP(&internal.Args.LoMethodsToExclude, "exclude", "x", []string{}, "exclude lo method e.g. --exclude=Map,Filter")
 
+	// include
+	rootCmd.Flags().StringSliceVarP(&internal.Args.RawLoMethodsToInclude, "include", "n", []string{}, "include lo method with regex e.g. --include=Filter*,Map")
+
 	// rename
 	rootCmd.Flags().StringSliceVarP(&internal.Args.RawRename, "rename", "r", []string{}, "rename method e.g. --rename=Map:Loop")
 }

example/user.go

@@ -94,5 +94,8 @@ type UserList []User
 //go:generate go run ../main.go --entity=User --slice=UserExclude --input=user.go --output=users_exclude_gen.go --exclude=Filter,KeyBy,GroupBy,FilterReject,Find
 type UserExclude []User
 
+//go:generate go run ../main.go --entity=User --slice=UserInclude --input=user.go --output=users_include_gen.go --include=Filter$,FilterBy*,Map,Key
+type UserInclude []User
+
 //go:generate go run ../main.go --entity=User --slice=UserRename --input=user.go --output=users_rename_gen.go --rename=Map:Loop --rename=Filter:LoFilter --rename=KeyBy:LoKeyBy
 type UserRename []User

internal/args.go

@@ -2,6 +2,8 @@ package internal
 
 import (
 	"fmt"
+	"regexp"
+	"slices"
 	"strings"
 )
 
@@ -32,21 +34,36 @@ type Arguments struct {
 
 	// Excluded lo methods
 	LoMethodsToExclude []string
+
+	// Included lo methods
+	RawLoMethodsToInclude []string
+	LoMethodsToInclude    []regexp.Regexp
 }
 
 var Args = Arguments{
 	RenameMap: map[string]string{},
 }
 
 func (a *Arguments) Load() error {
+	container := make([]error, 0)
+
 	// Load rename
 	if err := a.loadRename(a.RawRename); err != nil {
-		return fmt.Errorf("load accessor error: %w", err)
+		container = append(container, fmt.Errorf("load rename error: %w", err))
 	}
 
 	// Load entity
 	if err := a.loadEntity(a.RawEntity); err != nil {
-		return fmt.Errorf("load entity error: %w", err)
+		container = append(container, fmt.Errorf("load entity error: %w", err))
+	}
+
+	// Load include
+	if err := a.loadLoMethodsToInclude(); err != nil {
+		container = append(container, fmt.Errorf("load include error: %w", err))
+	}
+
+	if len(container) != 0 {
+		return fmt.Errorf("%v", container)
 	}
 
 	return nil
@@ -59,6 +76,26 @@ func (a Arguments) DisplayEntity() string {
 	return a.Entity
 }
 
+// Exclude
+// TODO: support regex
+func (a Arguments) IsExcluded(method string) bool {
+	return slices.Contains(a.LoMethodsToExclude, method)
+}
+
+// Include with regex
+func (a Arguments) IsIncluded(method string) bool {
+	if len(a.LoMethodsToInclude) == 0 { // default: include all
+		return true
+	}
+	for _, re := range a.LoMethodsToInclude {
+		if re.MatchString(method) {
+			return true
+		}
+	}
+	return false
+}
+
+// load Rename flag
 func (a *Arguments) loadRename(as []string) error {
 	container := make([]error, 0)
 	for _, ac := range as {
@@ -76,6 +113,7 @@ func (a *Arguments) loadRename(as []string) error {
 	return nil
 }
 
+// load Entity flag
 func (a *Arguments) loadEntity(e string) error {
 	a.IsPtrEntity = strings.HasPrefix(e, "*")
 	if a.IsPtrEntity {
@@ -84,3 +122,20 @@ func (a *Arguments) loadEntity(e string) error {
 	a.Entity = e
 	return nil
 }
+
+// load LoMethodsToInclude flag
+func (a *Arguments) loadLoMethodsToInclude() error {
+	container := make([]error, 0)
+	for _, raw := range a.RawLoMethodsToInclude {
+		re, err := regexp.Compile(raw)
+		if err != nil {
+			container = append(container, fmt.Errorf("invalid regex: %s", raw))
+			continue
+		}
+		a.LoMethodsToInclude = append(a.LoMethodsToInclude, *re)
+	}
+	if len(container) != 0 {
+		return fmt.Errorf("%v", container)
+	}
+	return nil
+}

internal/generator/generator.go

@@ -79,11 +79,16 @@ func (g Generator) genStd(args internal.Arguments, sliceName string) (string, er
 	var doc bytes.Buffer
 
 	for _, elem := range g.LoList {
-		// Skip if method should be excluded.
+		// Skip if method is excluded.
 		if slices.Contains(args.LoMethodsToExclude, elem.StdName()) {
 			continue
 		}
 
+		// Skip if method is NOT included.
+		if !args.IsIncluded(elem.StdName()) {
+			continue
+		}
+
 		// Get template src
 		rawTemp, ok := elem.StdTemplate()
 		if !ok {
@@ -123,11 +128,16 @@ func (g Generator) genExtend(args internal.Arguments, sliceName string, fields i
 	var doc bytes.Buffer
 
 	for _, elem := range g.LoList {
-		// Skip if method should be excluded.
+		// Skip if method is excluded.
 		if slices.Contains(args.LoMethodsToExclude, elem.StdName()) {
 			continue
 		}
 
+		// Skip if method is NOT included.
+		if !args.IsIncluded(elem.StdName()) {
+			continue
+		}
+
 		// Get template src
 		rawTemp, ok := elem.ExtendTemplate()
 		if !ok {