cmd,drawer: support dir walker

cmd/root.go

@@ -19,11 +19,13 @@ import (
 	"fmt"
 	"log"
 	"os"
+	"path/filepath"
 	"strings"
 
 	"github.com/scbizu/mew/drawer"
 	"github.com/scbizu/mew/filter"
 	"github.com/scbizu/mew/linker"
+	"github.com/scylladb/go-set/strset"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 )
@@ -35,28 +37,64 @@ var excludeDirs []string
 var isShowJSON bool
 var dumpGraph string
 var deepMode bool
+var dir string
+var dirEx []string
+var short bool
 
 // RootCmd represents the base command when called without any subcommands
 var RootCmd = &cobra.Command{
 	Use:   "mew",
 	Short: "mew - Show your Go repo related pkgs",
 	Long:  `mew - Show your Go repo related pkgs`,
-	Run: func(cmd *cobra.Command, args []string) {
-		l := linker.NewLinker(gopath, repoName)
-		if deepMode {
-			if jsonRes := handlePKGMap(l, excludeDirs, grep, dumpGraph, repoName, isShowJSON); jsonRes != "" {
-				fmt.Println(jsonRes)
+	RunE: func(cmd *cobra.Command, args []string) error {
+		repoNames := []string{repoName}
+		if dir != "" {
+			repoNames = []string{}
+			s := strset.New(dirEx...)
+			// ignore the repoName
+			if err := filepath.Walk(dir, func(path string, info os.FileInfo, _ error) error {
+				if info.IsDir() && !s.Has(info.Name()) {
+					repoNames = append(repoNames, filepath.Clean(strings.TrimPrefix(path, gopath+"/src/")))
+				}
+				return nil
+			}); err != nil {
+				return err
 			}
-		} else {
-			if jsonRes := handlePKGSlice(l, excludeDirs, grep, dumpGraph, repoName, isShowJSON); jsonRes != "" {
-				fmt.Println(jsonRes)
+		}
+		d := drawer.NewDot()
+		for _, repoName := range repoNames {
+			l := linker.NewLinker(gopath, repoName)
+			if deepMode {
+				if jsonRes := handlePKGMap(l, excludeDirs, grep, dumpGraph, repoName, isShowJSON); jsonRes != "" {
+					fmt.Println(jsonRes)
+				}
+			} else {
+				pkgs, err := getpkgs(l, excludeDirs, grep, dumpGraph, repoName)
+				if err != nil {
+					return err
+				}
+				repoName = shortFunc()(repoName)
+				if err := d.AddDep(repoName, pkgs); err != nil {
+					return err
+				}
 			}
 		}
-
-		return
+		if err := d.WriteFile(dumpGraph); err != nil {
+			return err
+		}
+		return nil
 	},
 }
 
+func shortFunc() func(string) string {
+	return func(s string) string {
+		if short {
+			return filepath.Base(s)
+		}
+		return s
+	}
+}
+
 func handlePKGMap(l *linker.Linker, excludeDirs []string, grep string, graphName string, repo string, isShowJSON bool) string {
 	var pkgMap map[string][]string
 	var err error
@@ -82,28 +120,19 @@ func handlePKGMap(l *linker.Linker, excludeDirs []string, grep string, graphName
 	return ""
 }
 
-func handlePKGSlice(l *linker.Linker, excludeDirs []string, grep string, graphName string, repo string, isShowJSON bool) string {
+func getpkgs(l *linker.Linker, excludeDirs []string, grep string, graphName string, repo string) ([]string, error) {
 	var pkgs []string
 	var err error
 	pkgs, err = l.GetLayerPKGNames(false, excludeDirs)
 	if err != nil {
-		logrus.Fatalln(err)
+		return nil, err
 	}
 	pkgFilter := filter.NewFilter(pkgs)
 	pkgs = pkgFilter.Grep(grep)
-
-	if err = drawer.DrawWithSliceAndSave(graphName, repo, pkgs); err != nil {
-		logrus.Fatalln(err.Error())
+	for index := range pkgs {
+		pkgs[index] = shortFunc()(pkgs[index])
 	}
-
-	if isShowJSON {
-		jsonRes, err := json.Marshal(pkgs)
-		if err != nil {
-			logrus.Fatalln(err)
-		}
-		return string(jsonRes)
-	}
-	return ""
+	return pkgs, nil
 }
 
 // Execute adds all child commands to the root command sets flags appropriately.
@@ -125,8 +154,11 @@ func init() {
 	}
 	RootCmd.Flags().StringVarP(&repoName, "repo", "r", "", "input repo name")
 	RootCmd.Flags().StringVarP(&grep, "grep", "g", "", "grep the pkg list")
+	RootCmd.Flags().StringVar(&dir, "dir", os.Getenv("GOPATH"), "the whole directory")
 	RootCmd.Flags().StringArrayVarP(&excludeDirs, "ed", "e", []string{"vendor", ".git"}, "exclude the dir")
+	RootCmd.Flags().StringArrayVar(&dirEx, "ex", []string{"vendor"}, "exclude the dir")
 	RootCmd.Flags().BoolVar(&isShowJSON, "json", false, "show json format")
+	RootCmd.Flags().BoolVar(&short, "short", false, "file base name")
 	RootCmd.Flags().BoolVar(&deepMode, "deep", false, "[Experimental feature]in deep mode,you will get all(include really all dependency) third party related pkg name")
 	RootCmd.Flags().StringVarP(&dumpGraph, "graph", "d", drawer.DefaultFileName, "dump graphviz graph")
 }

drawer/drawer.go

@@ -2,6 +2,7 @@
 package drawer
 
 import (
+	"bytes"
 	"io/ioutil"
 	"os"
 	"strconv"
@@ -18,10 +19,7 @@ const (
 )
 
 // DrawWithSlice returns the DOT lang  of a slice
-func DrawWithSlice(baseNode string, pkgNames []string) (string, error) {
-	g := gographviz.NewGraph()
-	_ = g.SetName(GraphName)
-	_ = g.SetDir(true)
+func DrawWithSlice(g *gographviz.Graph, baseNode string, pkgNames []string) (string, error) {
 	paresedNodeName := addQuotation(baseNode)
 	_ = g.AddNode(GraphName, paresedNodeName, nil)
 	for _, name := range pkgNames {
@@ -31,9 +29,45 @@ func DrawWithSlice(baseNode string, pkgNames []string) (string, error) {
 	return g.String(), nil
 }
 
+type DotTree struct {
+	g      *gographviz.Graph
+	buffer *bytes.Buffer
+}
+
+func NewDot() *DotTree {
+	g := gographviz.NewGraph()
+	g.SetName(GraphName)
+	g.SetDir(true)
+	return &DotTree{g: g, buffer: bytes.NewBufferString(g.String())}
+}
+
+func (d *DotTree) AddDep(baseNode string, pkgNames []string) error {
+	paresedNodeName := addQuotation(baseNode)
+	_ = d.g.AddNode(GraphName, paresedNodeName, nil)
+	for _, name := range pkgNames {
+		_ = d.g.AddNode(GraphName, addQuotation(name), nil)
+		_ = d.g.AddEdge(paresedNodeName, addQuotation(name), true, nil)
+	}
+	d.buffer.Reset()
+	if _, err := d.buffer.WriteString(d.g.String()); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (d *DotTree) WriteFile(filename string) error {
+	if err := ioutil.WriteFile(filename, d.buffer.Bytes(), os.ModePerm); err != nil {
+		return err
+	}
+	return nil
+}
+
 // DrawWithSliceAndSave draws the grpah and save to current path by a slice
 func DrawWithSliceAndSave(filename string, baseNode string, pkgNames []string) error {
-	dotTree, err := DrawWithSlice(baseNode, pkgNames)
+	g := gographviz.NewGraph()
+	_ = g.SetName(GraphName)
+	_ = g.SetDir(true)
+	dotTree, err := DrawWithSlice(g, baseNode, pkgNames)
 	if err != nil {
 		return err
 	}