output in tree form

internal/cmd/node.go

@@ -25,9 +25,10 @@ func init() {
 }
 
 var nodeCmd = &cobra.Command{
-	Args:  cobra.NoArgs,
-	Use:   "node",
-	Short: "Provides insight into the distribution of nodes per region or zone.",
+	Args:    cobra.NoArgs,
+	Use:     "nodes",
+	Aliases: []string{"node"},
+	Short:   "Provides insight into the distribution of nodes per region or zone.",
 	RunE: func(cmd *cobra.Command, args []string) error {
 		r, _ := cmd.Flags().GetString("region")
 		z, _ := cmd.Flags().GetString("zone")
@@ -39,6 +40,13 @@ var nodeCmd = &cobra.Command{
 		if err != nil {
 			return err
 		}
+		t, err := cmd.Flags().GetBool("tree")
+		if err != nil {
+			return err
+		}
+		if t {
+			return util.PrintNodeTree(n)
+		}
 		return util.PrintResult(n, false)
 	},
 }

internal/cmd/pod.go

@@ -27,9 +27,10 @@ func init() {
 }
 
 var podCmd = &cobra.Command{
-	Args:  cobra.NoArgs,
-	Use:   "pod",
-	Short: "Provides insight into the distribution of pods per region or zone.",
+	Args:    cobra.NoArgs,
+	Use:     "pods",
+	Aliases: []string{"pod"},
+	Short:   "Provides insight into the distribution of pods per region or zone.",
 	RunE: func(cmd *cobra.Command, args []string) error {
 		a, err := cmd.Flags().GetBool("all-namespaces")
 		if err != nil {
@@ -49,6 +50,13 @@ var podCmd = &cobra.Command{
 		if err != nil {
 			return err
 		}
+		t, err := cmd.Flags().GetBool("tree")
+		if err != nil {
+			return err
+		}
+		if t {
+			return util.PrintPodTree(p)
+		}
 		return util.PrintResult(p, false)
 	},
 }

internal/cmd/root.go

@@ -30,6 +30,7 @@ func init() {
 	configFlags.AddFlags(rootCmd.PersistentFlags())
 	rootCmd.PersistentFlags().StringP("region", "r", "", "The region to filter resources by. Mutually exclusive with '--zone'.")
 	rootCmd.PersistentFlags().StringP("zone", "z", "", "The zone to filter resources by. Mutually exclusive with '--region'.")
+	rootCmd.PersistentFlags().Bool("tree", false, "Print output in tree form.")
 	rootCmd.SetVersionTemplate("kubectl-topology " + version.Version)
 }
 

internal/util/tree.go

@@ -0,0 +1,148 @@
+package util
+
+import (
+	"fmt"
+	"sort"
+	"strings"
+)
+
+const (
+	firstElemPrefix  = `├─`
+	middleElemPrefix = `│ `
+	lastElemPrefix   = `└─`
+	indent           = "  "
+)
+
+type treeElement struct {
+	name string
+	subs []*treeElement
+}
+
+func newTreeElement(name string) *treeElement {
+	return &treeElement{
+		name: name,
+		subs: make([]*treeElement, 0),
+	}
+}
+
+func (te *treeElement) addSub(sub *treeElement) {
+	te.subs = append(te.subs, sub)
+	sort.SliceStable(te.subs, func(i, j int) bool {
+		return te.subs[i].name < te.subs[j].name
+	})
+}
+
+func (te *treeElement) treeString(linePrefix string) string {
+	sb := strings.Builder{}
+	sb.WriteString(fmt.Sprintf("%s\n", te.name))
+	if te.subs == nil || len(te.subs) == 0 {
+		return sb.String()
+	}
+	subSb := strings.Builder{}
+	for i, sub := range te.subs {
+		last := i == len(te.subs)-1
+		var subLinePrefix string
+		if !last {
+			subSb.WriteString(firstElemPrefix)
+			subLinePrefix = middleElemPrefix
+		} else {
+			subSb.WriteString(lastElemPrefix)
+			subLinePrefix = indent
+		}
+		subSb.WriteString(sub.treeString(subLinePrefix))
+	}
+	sb.WriteString(linesPrefix(subSb.String(), linePrefix))
+	return sb.String()
+}
+
+func PrintNodeTree(nodes []Node) error {
+	regionMap := groupNodes(nodes, func(node Node) string {
+		return node.Region
+	})
+	te := newTreeElement("cluster")
+	for region, nodes := range regionMap {
+		zoneMap := groupNodes(nodes, func(node Node) string {
+			return node.Zone
+		})
+		regionTE := newTreeElement(region)
+		te.addSub(regionTE)
+		for zone, nodes := range zoneMap {
+			zoneTE := newTreeElement(zone)
+			regionTE.addSub(zoneTE)
+			for _, node := range nodes {
+				zoneTE.addSub(newTreeElement(node.Name))
+			}
+		}
+	}
+	treeStr := te.treeString("")
+	fmt.Println(treeStr)
+	return nil
+}
+
+func PrintPodTree(pods []Pod) error {
+	regionMap := groupPods(pods, func(pod Pod) string {
+		return pod.Node.Region
+	})
+	te := newTreeElement("cluster")
+	for region, pods := range regionMap {
+		zoneMap := groupPods(pods, func(pod Pod) string {
+			return pod.Node.Zone
+		})
+		regionTE := newTreeElement(region)
+		te.addSub(regionTE)
+		for zone, pods := range zoneMap {
+			nodeMap := groupPods(pods, func(pod Pod) string {
+				return pod.Node.Name
+			})
+			zoneTE := newTreeElement(zone)
+			regionTE.addSub(zoneTE)
+			for node, pods := range nodeMap {
+				nodeTE := newTreeElement(node)
+				zoneTE.addSub(nodeTE)
+				for _, pod := range pods {
+					nodeTE.addSub(newTreeElement(fmt.Sprintf("%s - %s", pod.Namespace, pod.Name)))
+				}
+			}
+		}
+	}
+	treeStr := te.treeString("")
+	fmt.Println(treeStr)
+	return nil
+}
+
+func groupNodes(nodes []Node, keyFunc func(Node) string) map[string][]Node {
+	m := make(map[string][]Node)
+	for _, node := range nodes {
+		key := keyFunc(node)
+		n, ok := m[key]
+		if !ok {
+			n = make([]Node, 0)
+		}
+		m[key] = append(n, node)
+	}
+	return m
+}
+
+func groupPods(pods []Pod, keyFunc func(Pod) string) map[string][]Pod {
+	m := make(map[string][]Pod)
+	for _, pod := range pods {
+		key := keyFunc(pod)
+		n, ok := m[key]
+		if !ok {
+			n = make([]Pod, 0)
+		}
+		m[key] = append(n, pod)
+	}
+	return m
+}
+
+func linesPrefix(s string, prefix string) string {
+	lines := strings.Split(s, "\n")
+	sb := strings.Builder{}
+	for _, line := range lines {
+		if len(line) != 0 {
+			sb.WriteString(prefix + line + "\n")
+		}
+	}
+	return sb.String()
+}