diff --git a/cmd/main.go b/cmd/main.go index 9a7edfe..201d9ef 100644 --- a/cmd/main.go +++ b/cmd/main.go @@ -1,6 +1,7 @@ package main import ( + "flag" "fmt" "go/ast" "go/parser" @@ -8,14 +9,21 @@ import ( "log" "os" "path/filepath" + "strings" +) + +var ( + customLogs = flag.String("funcs", "", "Custom functions to export delimited by ','.") ) func main() { - if len(os.Args) != 2 { + flag.Parse() + + if flag.NArg() != 1 { log.Fatalln("directory must be provided") } - directories, err := getAllSubdirectories(os.Args[1]) + directories, err := getAllSubdirectories(flag.Arg(0)) if err != nil { log.Fatalf("failed to list subdirs: %v", err) } @@ -78,7 +86,27 @@ func isLogDot(expr ast.Expr) bool { return ok && (isIdent(sel.Sel, "Debug") || isIdent(sel.Sel, "Info") || isIdent(sel.Sel, "Warn") || - isIdent(sel.Sel, "Error")) + isIdent(sel.Sel, "Error") || + isCustomLog(sel.Sel)) +} + +func isCustomLog(expr ast.Expr) bool { + id, ok := expr.(*ast.Ident) + if !ok { + return false + } + + if len(*customLogs) == 0 { + return false + } + + ss := strings.Split(*customLogs, ",") + for i := range ss { + if id.Name == ss[i] { + return true + } + } + return false } func isIdent(expr ast.Expr, ident string) bool {