From 01b5c0e3d93caec5067c745cfb191952092d4719 Mon Sep 17 00:00:00 2001 From: Antoine Gelloz Date: Wed, 21 Aug 2024 16:32:57 +0200 Subject: [PATCH] feat: add graph printer --- .gitignore | 2 +- Makefile | 6 ++- pkg/dag/dag.go | 24 ++++++++++++ pkg/dag/dag_test.go | 70 +++++++++++++++++++++++++++++++++ pkg/dag/printer.go | 95 +++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 195 insertions(+), 2 deletions(-) create mode 100644 pkg/dag/printer.go diff --git a/.gitignore b/.gitignore index 07268668..ee3f3b67 100644 --- a/.gitignore +++ b/.gitignore @@ -9,6 +9,6 @@ docs/cmd pkg/dib/tests pkg/trivy/reports venv -coverage.output +coverage.* .golangci.yml golangci-lint-report.xml diff --git a/Makefile b/Makefile index cbf82802..0a075bba 100644 --- a/Makefile +++ b/Makefile @@ -42,9 +42,13 @@ RESET := $(shell tput sgr0) .PHONY: test test: ## Run tests - @go test -v -race -failfast -coverprofile coverage.output -run $(RUN) $(PKG) | \ + @go test -v -race -failfast -coverprofile coverage.out -covermode atomic -run $(RUN) $(PKG) | \ sed 's/RUN/$(BLUE)RUN$(RESET)/g' | \ sed 's/CONT/$(BLUE)CONT$(RESET)/g' | \ sed 's/PAUSE/$(BLUE)PAUSE$(RESET)/g' | \ sed 's/PASS/$(GREEN)PASS$(RESET)/g' | \ sed 's/FAIL/$(RED)FAIL$(RESET)/g' + @go tool cover -html=coverage.out -o coverage.html + @echo "To open the html coverage file, use one of the following commands:" + @echo "open coverage.html on mac" + @echo "xdg-open coverage.html on linux" diff --git a/pkg/dag/dag.go b/pkg/dag/dag.go index 15cb8725..c96ff31c 100644 --- a/pkg/dag/dag.go +++ b/pkg/dag/dag.go @@ -1,6 +1,9 @@ package dag import ( + "cmp" + "slices" + "strings" "sync" "golang.org/x/sync/errgroup" @@ -204,3 +207,24 @@ func createUniqueVisitorErr(visitor NodeVisitorFuncErr) NodeVisitorFuncErr { return uniqueVisitor } + +func sort(a, b *Node) int { + return cmp.Compare( + strings.ToLower(a.Image.ShortName), + strings.ToLower(b.Image.ShortName), + ) +} + +func (d *DAG) Sprint(name string) string { + d.WalkInDepth(func(node *Node) { + slices.SortFunc(node.Children(), sort) + }) + + slices.SortFunc(d.nodes, sort) + rootNode := &Node{ + Image: &Image{Name: name}, + children: d.nodes, + } + + return defaultPrinter.WithRoot(rootNode).Srender() +} diff --git a/pkg/dag/dag_test.go b/pkg/dag/dag_test.go index b7c3309d..15f93c68 100644 --- a/pkg/dag/dag_test.go +++ b/pkg/dag/dag_test.go @@ -271,3 +271,73 @@ func Test_ListImage(t *testing.T) { " hash: arm-ag-ed-don\n" assert.Equal(t, expected, DAG.ListImage()) } + +func TestDAG_Sprint(t *testing.T) { + t.Parallel() + + t.Run("empty graph", func(t *testing.T) { + t.Parallel() + + graph := &dag.DAG{} + node := dag.NewNode(nil) + graph.AddNode(node) + + assert.Equal(t, "graph\n", graph.Sprint("graph")) + }) + + t.Run("one node", func(t *testing.T) { + t.Parallel() + + graph := &dag.DAG{} + node := dag.NewNode(&dag.Image{ + ShortName: "n1", + Hash: "h1", + }) + graph.AddNode(node) + + expected := `graph +└───n1 [h1] +` + assert.Equal(t, expected, graph.Sprint("graph")) + }) + + t.Run("5 nodes", func(t *testing.T) { + t.Parallel() + + graph := &dag.DAG{} + node1 := dag.NewNode(&dag.Image{ + ShortName: "n1", + Hash: "h1", + }) + node2 := dag.NewNode(&dag.Image{ + ShortName: "n2", + Hash: "h2", + }) + node3 := dag.NewNode(&dag.Image{ + ShortName: "n3", + Hash: "h3", + }) + node4 := dag.NewNode(&dag.Image{ + ShortName: "n4", + Hash: "h4", + }) + node5 := dag.NewNode(&dag.Image{ + ShortName: "n5", + Hash: "h5", + }) + node1.AddChild(node2) + node1.AddChild(node3) + node1.AddChild(node4) + node2.AddChild(node5) + graph.AddNode(node1) + + expected := `graph +└──┬n1 [h1] + ├──┬n2 [h2] + │ └───n5 [h5] + ├───n3 [h3] + └───n4 [h4] +` + assert.Equal(t, expected, graph.Sprint("graph")) + }) +} diff --git a/pkg/dag/printer.go b/pkg/dag/printer.go new file mode 100644 index 00000000..06829d3a --- /dev/null +++ b/pkg/dag/printer.go @@ -0,0 +1,95 @@ +package dag + +import ( + "fmt" + "io" + "strings" + + "github.com/pterm/pterm" +) + +var defaultPrinter = GraphPrinter{ + TopRightCornerString: "└", + TopRightDownString: "├", + HorizontalString: "─", + VerticalString: "│", + RightDownLeftString: "┬", + Indent: 3, +} + +type GraphPrinter struct { + Root *Node + TreeStyle *pterm.Style + TextStyle *pterm.Style + TopRightCornerString string + TopRightDownString string + HorizontalString string + VerticalString string + RightDownLeftString string + Indent int + Writer io.Writer +} + +// WithRoot returns a new GraphPrinter with a specific Root node. +func (p GraphPrinter) WithRoot(root *Node) *GraphPrinter { + p.Root = root + return &p +} + +// Srender renders the graph as a string. +func (p GraphPrinter) Srender() string { + if p.TreeStyle == nil { + p.TreeStyle = pterm.NewStyle() + } + if p.TextStyle == nil { + p.TextStyle = pterm.NewStyle() + } + + var result string + if p.Root.Image.Name != "" { + result += p.TextStyle.Sprint(p.Root.Image.Name) + "\n" + } + result += walkOverTree(p.Root.Children(), p, "") + + return result +} + +func walkOverTree(nodes []*Node, printer GraphPrinter, prefix string) string { + var result string + for nodeIndex, node := range nodes { + if node.Image == nil { + continue + } + + txt := fmt.Sprintf("%s [%s]\n", node.Image.ShortName, node.Image.Hash) + if len(nodes) > nodeIndex+1 { // if not last in nodes + if len(node.Children()) == 0 { // if there are no children + result += prefix + printer.TreeStyle.Sprint(printer.TopRightDownString) + + strings.Repeat(printer.TreeStyle.Sprint(printer.HorizontalString), printer.Indent) + + printer.TextStyle.Sprint(txt) + } else { // if there are children + result += prefix + printer.TreeStyle.Sprint(printer.TopRightDownString) + + strings.Repeat(printer.TreeStyle.Sprint(printer.HorizontalString), printer.Indent-1) + + printer.TreeStyle.Sprint(printer.RightDownLeftString) + + printer.TextStyle.Sprint(txt) + result += walkOverTree(node.Children(), printer, + prefix+printer.TreeStyle.Sprint(printer.VerticalString)+strings.Repeat(" ", printer.Indent-1)) + } + } else if len(nodes) == nodeIndex+1 { // if last in nodes + if len(node.Children()) == 0 { // if there are no children + result += prefix + printer.TreeStyle.Sprint(printer.TopRightCornerString) + + strings.Repeat(printer.TreeStyle.Sprint(printer.HorizontalString), printer.Indent) + + printer.TextStyle.Sprint(txt) + } else { // if there are children + result += prefix + printer.TreeStyle.Sprint(printer.TopRightCornerString) + + strings.Repeat(printer.TreeStyle.Sprint(printer.HorizontalString), printer.Indent-1) + + printer.TreeStyle.Sprint(printer.RightDownLeftString) + + printer.TextStyle.Sprint(txt) + result += walkOverTree(node.Children(), printer, + prefix+strings.Repeat(" ", printer.Indent)) + } + } + } + + return result +}