improve display logic (#68)

* improve display logic

* add test

internal/check/check.go

@@ -35,7 +35,7 @@ func Check[T any](parser dep_tree.NodeParser[T], cfg *Config) error {
 			formattedCycleStack := make([]string, len(el.Value.Stack))
 			for i, el := range el.Value.Stack {
 				if node := dt.Graph.Get(el); node != nil {
-					formattedCycleStack[i] = parser.Display(node)
+					formattedCycleStack[i] = parser.Display(node).Name
 				} else {
 					formattedCycleStack[i] = el
 				}

internal/dep_tree/dep_tree.go

@@ -13,8 +13,13 @@ import (
 
 type NodeParserBuilder[T any] func([]string) (NodeParser[T], error)
 
+type DisplayResult struct {
+	Name  string
+	Group string
+}
+
 type NodeParser[T any] interface {
-	Display(node *graph.Node[T]) string
+	Display(node *graph.Node[T]) DisplayResult
 	Node(id string) (*graph.Node[T], error)
 	Deps(node *graph.Node[T]) ([]*graph.Node[T], error)
 }
@@ -77,13 +82,13 @@ func (dt *DepTree[T]) WithStdErrLoader() *DepTree[T] {
 	dt.onNodeStartLoad = func(n *graph.Node[T]) {
 		done += 1
 		_ = bar.Set(done)
-		bar.Describe(fmt.Sprintf("(%d/%d) Loading %s...", done, len(diff), dt.NodeParser.Display(n)))
+		bar.Describe(fmt.Sprintf("(%d/%d) Loading %s...", done, len(diff), dt.NodeParser.Display(n).Name))
 	}
 	dt.onNodeFinishLoad = func(n *graph.Node[T], ns []*graph.Node[T]) {
 		for _, n := range ns {
 			diff[n.Id] = true
 		}
-		bar.Describe(fmt.Sprintf("(%d/%d) Loading %s...", done, len(diff), dt.NodeParser.Display(n)))
+		bar.Describe(fmt.Sprintf("(%d/%d) Loading %s...", done, len(diff), dt.NodeParser.Display(n).Name))
 	}
 	dt.onFinishLoad = func() {
 		bar.Describe("Finished loading")

internal/dep_tree/render.go

@@ -53,7 +53,7 @@ func (dt *DepTree[T]) Render() (*board.Board, error) {
 		err := b.AddBlock(
 			&board.Block{
 				Id:       n.Node.Id,
-				Label:    prefix + dt.NodeParser.Display(n.Node),
+				Label:    prefix + dt.NodeParser.Display(n.Node).Name,
 				Position: utils.Vec(indent*n.Lvl+xOffset, i+yOffset),
 				Tags:     tags,
 			},

internal/dep_tree/structured.go

@@ -31,7 +31,7 @@ func (dt *DepTree[T]) makeStructuredTree(
 			result = make(map[string]interface{})
 		}
 		var err error
-		result[dt.NodeParser.Display(to)], err = dt.makeStructuredTree(to.Id, nil)
+		result[dt.NodeParser.Display(to).Name], err = dt.makeStructuredTree(to.Id, nil)
 		if err != nil {
 			return nil, err
 		}
@@ -52,7 +52,7 @@ func (dt *DepTree[T]) RenderStructured() ([]byte, error) {
 
 	structuredTree := StructuredTree{
 		Tree: map[string]interface{}{
-			dt.NodeParser.Display(dt.Entrypoints[0]): tree,
+			dt.NodeParser.Display(dt.Entrypoints[0]).Name: tree,
 		},
 		CircularDependencies: make([][]string, 0),
 		Errors:               make(map[string][]string),
@@ -62,14 +62,14 @@ func (dt *DepTree[T]) RenderStructured() ([]byte, error) {
 		cycleDep, _ := dt.Cycles.Get(cycle)
 		renderedCycle := make([]string, len(cycleDep.Stack))
 		for i, cycleDepEntry := range cycleDep.Stack {
-			renderedCycle[i] = dt.NodeParser.Display(dt.Graph.Get(cycleDepEntry))
+			renderedCycle[i] = dt.NodeParser.Display(dt.Graph.Get(cycleDepEntry)).Name
 		}
 		structuredTree.CircularDependencies = append(structuredTree.CircularDependencies, renderedCycle)
 	}
 
 	for _, node := range dt.Nodes {
 		if node.Node.Errors != nil && len(node.Node.Errors) > 0 {
-			erroredNode := dt.NodeParser.Display(dt.Graph.Get(node.Node.Id))
+			erroredNode := dt.NodeParser.Display(dt.Graph.Get(node.Node.Id)).Name
 			nodeErrors := make([]string, len(node.Node.Errors))
 			for i, err := range node.Node.Errors {
 				nodeErrors[i] = err.Error()

internal/dep_tree/test_parser.go

@@ -42,6 +42,6 @@ func (t *TestParser) Deps(n *graph.Node[[]int]) ([]*graph.Node[[]int], error) {
 	return result, nil
 }
 
-func (t *TestParser) Display(n *graph.Node[[]int]) string {
-	return n.Id
+func (t *TestParser) Display(n *graph.Node[[]int]) DisplayResult {
+	return DisplayResult{Name: n.Id}
 }

internal/entropy/dirs.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 
 	"github.com/elliotchance/orderedmap/v2"
+	"github.com/gabotechs/dep-tree/internal/language"
 
 	"github.com/gabotechs/dep-tree/internal/utils"
 )
@@ -52,47 +53,54 @@ func splitBaseNames(dir string) []string {
 	return result
 }
 
-func (d *DirTree) AddDirs(dir string) {
+func (d *DirTree) AddDirs(dirs []string) {
 	node := d.inner()
-	for _, p := range splitBaseNames(dir) {
-		base := filepath.Base(p)
-		if upper, ok := node.Get(base); ok {
+	for _, dir := range dirs {
+		if upper, ok := node.Get(dir); ok {
 			node = upper.entry.inner()
 		} else {
 			newNode := NewDirTree()
-			node.Set(base, DirTreeEntry{newNode, node.Len()})
+			node.Set(dir, DirTreeEntry{newNode, node.Len()})
 			node = newNode.inner()
 		}
 	}
 }
 
-// ColorFor smartly assigns a color for the specified dir based on all the dir tree that
+func (d *DirTree) AddDirsFromDisplay(display language.DisplayResult) {
+	dirs := splitBaseNames(filepath.Dir(display.Name))
+	if display.Group != "" {
+		d.AddDirs(utils.AppendFront(display.Group, dirs))
+	} else {
+		d.AddDirs(dirs)
+	}
+}
+
+// ColorForDir smartly assigns a color for the specified dir based on all the dir tree that
 // the codebase has. Files in the same folder will receive the same color, and colors for
 // each sub folder will be assigned evenly following a radial distribution in an HSV wheel.
 // As it goes deeper into more sub folders, colors fade, but the distribution rules are
 // the same.
-func (d *DirTree) ColorFor(dir string) []int {
-	baseNames := splitBaseNames(dir)
+func (d *DirTree) ColorForDir(dirs []string) []int {
 	depth := 0
 	node := d.inner()
 	// It might happen that all the nodes have some common folders, like src/,
 	// so if literally all of them have the same common folders, we do not want to take
 	// them into account for distributing colors, as they will appear very faded.
-	for node.Len() == 1 && len(baseNames) > 0 {
-		if node.Front().Key != baseNames[0] {
+	for node.Len() == 1 && len(dirs) > 0 {
+		if node.Front().Key != dirs[0] {
 			break
 		}
 		node = node.Front().Value.entry.inner()
-		baseNames = baseNames[1:]
+		dirs = dirs[1:]
 	}
 	h, s, v := float64(0), 0., 1.
-	for depth < len(baseNames) {
-		el, ok := node.Get(baseNames[depth])
+	for depth < len(dirs) {
+		el, ok := node.Get(dirs[depth])
 		if !ok {
 			return []int{0, 0, 0}
 		}
 		h = float64(int(h+360*float64(el.index)/float64(node.Len())) % 360)
-		s = utils.Scale(1-float64(depth)/float64(len(baseNames)), 0, 1, .2, .9)
+		s = utils.Scale(1-float64(depth)/float64(len(dirs)), 0, 1, .2, .9)
 
 		depth += 1
 		node = el.entry.inner()
@@ -101,6 +109,46 @@ func (d *DirTree) ColorFor(dir string) []int {
 	return []int{int(r), int(g), int(b)}
 }
 
+func (d *DirTree) ColorForDisplay(display language.DisplayResult) []int {
+	dirs := splitBaseNames(filepath.Dir(display.Name))
+	if display.Group != "" {
+		return d.ColorForDir(utils.AppendFront(display.Group, dirs))
+	} else {
+		return d.ColorForDir(dirs)
+	}
+}
+
+func (d *DirTree) GroupingsForDir(dirs []string) []string {
+	depth := 0
+	var result []string
+
+	node := d.inner()
+	acc := ""
+	for depth < len(dirs) {
+		acc = filepath.Join(acc, dirs[depth])
+		el, ok := node.Get(dirs[depth])
+		if !ok {
+			return result
+		}
+		if node.Len() > 1 {
+			result = append(result, acc)
+		}
+
+		depth += 1
+		node = el.entry.inner()
+	}
+	return result
+}
+
+func (d *DirTree) GroupingsForDisplay(display language.DisplayResult) []string {
+	dirs := splitBaseNames(filepath.Dir(display.Name))
+	if display.Group != "" {
+		return d.GroupingsForDir(utils.AppendFront(display.Group, dirs))
+	} else {
+		return d.GroupingsForDir(dirs)
+	}
+}
+
 // HSVToRGB converts an HSV triple to an RGB triple.
 // taken from https://github.com/Crazy3lf/colorconv/blob/master/colorconv.go
 //

internal/entropy/dirs_test.go

@@ -6,7 +6,7 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func Test_dirs(t *testing.T) {
+func TestDirTree_baseFunctions(t *testing.T) {
 	tests := []struct {
 		Name              string
 		ExpectedFullPaths []string
@@ -101,14 +101,96 @@ func Test_dirs(t *testing.T) {
 		t.Run(tt.Name, func(t *testing.T) {
 			a := require.New(t)
 			dirTree := NewDirTree()
-			dirTree.AddDirs(tt.Name)
+			dirTree.AddDirs(splitBaseNames(tt.Name))
 			a.Equal(tt.ExpectedFullPaths, splitFullPaths(tt.Name))
 			a.Equal(tt.ExpectedBaseNames, splitBaseNames(tt.Name))
 			a.Equal(tt.ExpectedTree, unwrapDirTree(dirTree))
 		})
 	}
 }
 
+func TestDirTree_GroupingsForDir(t *testing.T) {
+	tests := []struct {
+		Name              string
+		Paths             []string
+		ExpectedTree      map[string]any
+		ExpectedGroupings [][]string
+	}{
+		{
+			Name:  "Single File",
+			Paths: []string{"foo/bar"},
+			ExpectedTree: map[string]any{
+				"foo": map[string]any{
+					"bar": map[string]any{},
+				},
+			},
+			ExpectedGroupings: [][]string{nil},
+		},
+		{
+			Name:  "two unrelated files",
+			Paths: []string{"foo/bar", "baz/bar"},
+			ExpectedTree: map[string]any{
+				"foo": map[string]any{
+					"bar": map[string]any{},
+				},
+				"baz": map[string]any{
+					"bar": map[string]any{},
+				},
+			},
+			ExpectedGroupings: [][]string{{"foo"}, {"baz"}},
+		},
+		{
+			Name:  "two files with a shared first folder",
+			Paths: []string{"foo/bar", "foo/baz"},
+			ExpectedTree: map[string]any{
+				"foo": map[string]any{
+					"bar": map[string]any{},
+					"baz": map[string]any{},
+				},
+			},
+			ExpectedGroupings: [][]string{{"foo/bar"}, {"foo/baz"}},
+		},
+		{
+			Name:  "with middle folders",
+			Paths: []string{"foo/bar/baz/1", "foo/bar/baz/2", "bar/foo"},
+			ExpectedTree: map[string]any{
+				"foo": map[string]any{
+					"bar": map[string]any{
+						"baz": map[string]any{
+							"1": map[string]any{},
+							"2": map[string]any{},
+						},
+					},
+				},
+				"bar": map[string]any{
+					"foo": map[string]any{},
+				},
+			},
+			ExpectedGroupings: [][]string{
+				{"foo", "foo/bar/baz/1"},
+				{"foo", "foo/bar/baz/2"},
+				{"bar"},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.Name, func(t *testing.T) {
+			a := require.New(t)
+			dirTree := NewDirTree()
+			for _, path := range tt.Paths {
+				dirTree.AddDirs(splitBaseNames(path))
+			}
+			a.Equal(tt.ExpectedTree, unwrapDirTree(dirTree))
+			var groupings [][]string
+			for _, path := range tt.Paths {
+				groupings = append(groupings, dirTree.GroupingsForDir(splitBaseNames(path)))
+			}
+			a.Equal(tt.ExpectedGroupings, groupings)
+		})
+	}
+}
+
 func unwrapDirTree(tree *DirTree) interface{} {
 	if tree == nil {
 		return nil

internal/entropy/graph.go

@@ -2,7 +2,6 @@ package entropy
 
 import (
 	"path/filepath"
-	"strings"
 
 	"github.com/gabotechs/dep-tree/internal/dep_tree"
 	"github.com/gabotechs/dep-tree/internal/graph"
@@ -17,6 +16,7 @@ const (
 type Node struct {
 	Id       int64  `json:"id"`
 	FileName string `json:"fileName"`
+	Group    string `json:"group,omitempty"`
 	DirName  string `json:"dirName"`
 	Loc      int    `json:"loc"`
 	Size     int    `json:"size"`
@@ -53,19 +53,20 @@ func makeGraph(dt *dep_tree.DepTree[language.FileInfo], parser language.NodePars
 	dirTree := NewDirTree()
 
 	for _, node := range allNodes {
-		dirTree.AddDirs(filepath.Dir(parser.Display(node)))
+		dirTree.AddDirsFromDisplay(parser.Display(node))
 	}
 
 	for _, node := range allNodes {
-		filePath := parser.Display(node)
-		dirName := filepath.Dir(filePath)
+		display := parser.Display(node)
+		dirName := filepath.Dir(display.Name)
 		out.Nodes = append(out.Nodes, Node{
 			Id:       node.ID(),
-			FileName: filepath.Base(filePath),
+			FileName: filepath.Base(display.Name),
+			Group:    display.Group,
 			DirName:  dirName + "/",
 			Loc:      node.Data.Loc,
 			Size:     maxNodeSize * node.Data.Loc / maxLoc,
-			Color:    dirTree.ColorFor(dirName),
+			Color:    dirTree.ColorForDisplay(display),
 		})
 
 		for _, to := range dt.Graph.FromId(node.Id) {
@@ -75,13 +76,7 @@ func makeGraph(dt *dep_tree.DepTree[language.FileInfo], parser language.NodePars
 			})
 		}
 
-		for _, parentFolder := range splitFullPaths(dirName) {
-			// NOTE: just ignore parent folders like ".." or "../..", otherwise they will contribute
-			//  to grouping folders that might be unrelated. Empirically, visualizations look nicer if
-			//  we ignore them.
-			if strings.HasSuffix(parentFolder, "..") {
-				continue
-			}
+		for _, parentFolder := range dirTree.GroupingsForDisplay(display) {
 			folderNode := graph.MakeNode(parentFolder, 0)
 			out.Links = append(out.Links, Link{
 				From:  node.ID(),