Merge pull request #9 from owlinux1000/colorized_output

Colorized output (Close #2)

README.md

@@ -35,6 +35,7 @@ Usage:
 
 Flags:
   -h, --help      help for gcstree
+  -n, --no-color   disable colorized outputs
   -v, --version   show the gcstree version
 ```
 

cmd/root.go

@@ -33,7 +33,12 @@ var rootCmd = &cobra.Command{
 		if err != nil {
 			log.Fatal(err)
 		}
-		gcsTree, err := internal.NewGCSTree(ctx, client, bucket)
+
+		withColorized, _ := cmd.Flags().GetBool("no-color")
+		option := internal.PrintOption{
+			WithColorized: !withColorized,
+		}
+		gcsTree, err := internal.NewGCSTree(ctx, client, bucket, &option)
 		if err != nil {
 			log.Fatal(err)
 		}
@@ -54,4 +59,5 @@ func Execute() {
 
 func init() {
 	rootCmd.Flags().BoolP("version", "v", false, "show the gcstree version")
+	rootCmd.Flags().BoolP("no-color", "n", false, "disable colorized outputs")
 }

go.mod

@@ -5,6 +5,7 @@ go 1.21.1
 require (
 	cloud.google.com/go/storage v1.33.0
 	github.com/ddddddO/gtree v1.9.11
+	github.com/fatih/color v1.15.0
 	github.com/spf13/cobra v1.7.0
 	google.golang.org/api v0.143.0
 )
@@ -14,7 +15,6 @@ require (
 	cloud.google.com/go/compute v1.23.0 // indirect
 	cloud.google.com/go/compute/metadata v0.2.3 // indirect
 	cloud.google.com/go/iam v1.1.1 // indirect
-	github.com/fatih/color v1.15.0 // indirect
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
 	github.com/golang/protobuf v1.5.3 // indirect
 	github.com/google/go-cmp v0.5.9 // indirect

internal/gcstree.go

@@ -16,9 +16,15 @@ type GCSTree struct {
 	folder  string
 	client  *storage.Client
 	counter *counter
+	option  *PrintOption
 }
 
-func NewGCSTree(ctx context.Context, client *storage.Client, bucket string) (*GCSTree, error) {
+type PrintOption struct {
+	WithColorized bool
+	WithSize      bool
+}
+
+func NewGCSTree(ctx context.Context, client *storage.Client, bucket string, option *PrintOption) (*GCSTree, error) {
 
 	folder := ""
 	if strings.Contains(bucket, "/") {
@@ -31,13 +37,15 @@ func NewGCSTree(ctx context.Context, client *storage.Client, bucket string) (*GC
 		folder:  folder,
 		client:  client,
 		counter: newCounter(),
+		option:  option,
 	}, nil
 }
 
 func (g *GCSTree) GetObjectAttrList(ctx context.Context) ([]*storage.ObjectAttrs, error) {
 	bkt := g.client.Bucket(g.bucket)
 	query := &storage.Query{Prefix: g.folder}
-	var names []*storage.ObjectAttrs
+	var objectAttrs []*storage.ObjectAttrs
+
 	it := bkt.Objects(ctx, query)
 	for {
 		attrs, err := it.Next()
@@ -48,9 +56,9 @@ func (g *GCSTree) GetObjectAttrList(ctx context.Context) ([]*storage.ObjectAttrs
 			return nil, err
 		}
 		g.counter.count(attrs.Name)
-		names = append(names, attrs)
+		objectAttrs = append(objectAttrs, attrs)
 	}
-	return names, nil
+	return objectAttrs, nil
 }
 
 func (g *GCSTree) String() (string, error) {
@@ -59,7 +67,7 @@ func (g *GCSTree) String() (string, error) {
 	if err != nil {
 		return "", err
 	}
-	treeResult, err := tree(g.bucket, objList)
+	treeResult, err := tree(g.bucket, objList, g.option)
 	if err != nil {
 		return "", err
 	}

internal/tree.go

@@ -1,21 +1,32 @@
 package internal
 
 import (
+	"path/filepath"
 	"strings"
 
 	"cloud.google.com/go/storage"
 	"github.com/ddddddO/gtree"
+	"github.com/fatih/color"
 )
 
 // ref: https://github.com/ddddddO/gtree#the-program-below-converts-the-result-of-find-into-a-tree
-func tree(bucket string, objList []*storage.ObjectAttrs) (string, error) {
+func tree(bucket string, objList []*storage.ObjectAttrs, option *PrintOption) (string, error) {
+	if option.WithColorized {
+		bucket = color.CyanString("%s", bucket)
+	}
 	root := gtree.NewRoot(bucket)
 	node := root
 	for _, obj := range objList {
+		_, file := filepath.Split(obj.Name)
 		for _, s := range strings.Split(obj.Name, "/") {
 			if s == "" {
 				continue
 			}
+			if option.WithColorized {
+				if s != file {
+					s = color.BlueString("%s", s)
+				}
+			}
 			tmp := node.Add(s)
 			node = tmp
 		}

internal/tree_test.go

@@ -21,7 +21,11 @@ func TestTree(t *testing.T) {
 		Size: 6,
 	})
 	bucket := "test"
-	got, err := tree(bucket, objList)
+	option := PrintOption{
+		WithColorized: false,
+		WithSize:      false,
+	}
+	got, err := tree(bucket, objList, &option)
 	if err != nil {
 		t.Fatal(err)
 	}