Skip to content

Commit

Permalink
feat: add graph printer
Browse files Browse the repository at this point in the history
  • Loading branch information
Antoine Gelloz committed Aug 21, 2024
1 parent 823adf4 commit 01b5c0e
Show file tree
Hide file tree
Showing 5 changed files with 195 additions and 2 deletions.
2 changes: 1 addition & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,6 @@ docs/cmd
pkg/dib/tests
pkg/trivy/reports
venv
coverage.output
coverage.*
.golangci.yml
golangci-lint-report.xml
6 changes: 5 additions & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -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"
24 changes: 24 additions & 0 deletions pkg/dag/dag.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
package dag

import (
"cmp"
"slices"
"strings"
"sync"

"golang.org/x/sync/errgroup"
Expand Down Expand Up @@ -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()
}
70 changes: 70 additions & 0 deletions pkg/dag/dag_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"))
})
}
95 changes: 95 additions & 0 deletions pkg/dag/printer.go
Original file line number Diff line number Diff line change
@@ -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
}

0 comments on commit 01b5c0e

Please sign in to comment.