drawer: support graphviz with deep mode

scbizu · May 26, 2018 · c27337f · c27337f
1 parent 6c88b97
commit c27337f
Show file tree

Hide file tree

Showing 4 changed files with 437 additions and 17 deletions.
diff --git a/cmd/root.go b/cmd/root.go
@@ -44,7 +44,7 @@ var RootCmd = &cobra.Command{
 	Run: func(cmd *cobra.Command, args []string) {
 		l := linker.NewLinker(gopath, repoName)
 		if deepMode {
-			if jsonRes := handlePKGMap(l, excludeDirs, grep, dumpGraph, isShowJSON); jsonRes != "" {
+			if jsonRes := handlePKGMap(l, excludeDirs, grep, dumpGraph, repoName, isShowJSON); jsonRes != "" {
 				fmt.Println(jsonRes)
 			}
 		} else {
@@ -57,7 +57,7 @@ var RootCmd = &cobra.Command{
 	},
 }
 
-func handlePKGMap(l *linker.Linker, excludeDirs []string, grep string, graphName string, isShowJSON bool) string {
+func handlePKGMap(l *linker.Linker, excludeDirs []string, grep string, graphName string, repo string, isShowJSON bool) string {
 	var pkgMap map[string][]string
 	var err error
 	pkgMap, err = l.GetAllPKGNames(false, excludeDirs)
@@ -68,7 +68,7 @@ func handlePKGMap(l *linker.Linker, excludeDirs []string, grep string, graphName
 	pkgMapFilter := filter.NewMapFilter(pkgMap)
 	pkgMapFilter.Grep(grep)
 
-	if err = drawer.DrawWithMapAndSave(graphName, pkgMap); err != nil {
+	if err = drawer.DrawWithMapAndSave(repo, graphName, pkgMap); err != nil {
 		logrus.Fatalln(err.Error())
 	}
 

diff --git a/drawer/drawer.go b/drawer/drawer.go
@@ -6,6 +6,7 @@ import (
 	"os"
 
 	"github.com/awalterschulze/gographviz"
+	"github.com/sirupsen/logrus"
 )
 
 const (
@@ -20,7 +21,7 @@ func DrawWithSlice(baseNode string, pkgNames []string) (string, error) {
 	g := gographviz.NewGraph()
 	g.SetName(GraphName)
 	g.SetDir(true)
-	paresedNodeName := string(parseBaseNodeName([]byte(baseNode)))
+	paresedNodeName := addQuotation(baseNode)
 	g.AddNode(GraphName, paresedNodeName, nil)
 	for _, name := range pkgNames {
 		g.AddNode(GraphName, name, nil)
@@ -42,23 +43,41 @@ func DrawWithSliceAndSave(filename string, baseNode string, pkgNames []string) e
 }
 
 // DrawWithMap returns DOT lang of a map
-func DrawWithMap(pkgMap map[string][]string) (string, error) {
+func DrawWithMap(baseNode string, pkgMap map[string][]string) (string, error) {
 	g := gographviz.NewGraph()
 	g.SetName(GraphName)
 	g.SetDir(true)
-	for repo, pkgs := range pkgMap {
-		paresedNodeName := string(parseBaseNodeName([]byte(repo)))
-		for _, pkg := range pkgs {
-			g.AddNode(GraphName, pkg, nil)
-			g.AddEdge(paresedNodeName, pkg, true, nil)
+
+	drawTree(g, baseNode, pkgMap)
+
+	return g.String(), nil
+}
+
+func drawTree(g *gographviz.Graph, base string, pkgMap map[string][]string) {
+
+	ps, ok := pkgMap[base]
+
+	if !ok || len(ps) == 0 {
+		return
+	}
+
+	g.AddNode(GraphName, addQuotation(base), nil)
+
+	for _, p := range ps {
+		g.AddNode(GraphName, addQuotation(p), nil)
+		if err := g.AddEdge(addQuotation(base), addQuotation(p), true, nil); err != nil {
+			logrus.Error(err)
 		}
+		drawTree(g, p, pkgMap)
 	}
-	return g.String(), nil
+
+	return
+
 }
 
 // DrawWithMapAndSave draws the graph and save to current path by a map
-func DrawWithMapAndSave(filename string, pkgMap map[string][]string) error {
-	dotTree, err := DrawWithMap(pkgMap)
+func DrawWithMapAndSave(baseNode string, filename string, pkgMap map[string][]string) error {
+	dotTree, err := DrawWithMap(baseNode, pkgMap)
 	if err != nil {
 		return err
 	}
@@ -68,11 +87,11 @@ func DrawWithMapAndSave(filename string, pkgMap map[string][]string) error {
 	return nil
 }
 
-func parseBaseNodeName(baseNode []byte) []byte {
+func addQuotation(node string) string {
 	bs := []byte{'"'}
-	for _, b := range baseNode {
+	for _, b := range []byte(node) {
 		bs = append(bs, b)
 	}
 	bs = append(bs, '"')
-	return bs
+	return string(bs)
 }