Merge pull request #47 from orangekame3/add-directory-only

[feat]: Add -d option

Author: orangekame3

README.md

@@ -251,6 +251,7 @@ Usage:
 
 Flags:
   -D, --date-time             Print the last modified time of each file.
+  -d, --directory-only        List directories only.
   -e, --endpoint-url string   AWS endpoint URL to use (useful for local testing with LocalStack)
   -f, --full-path             Print the full path prefix for each file.
   -h, --help                  help for stree
@@ -263,7 +264,9 @@ Flags:
   -p, --profile string        AWS profile to use (default "default")
   -r, --region string         AWS region to use (overrides the region specified in the profile)
   -s, --size                  Print the size of each file in bytes along with the name.
-  -u, --username              Print the owner of
+  -u, --username              Print the owner of each file.
+  -v, --version               version for stree
+
 
 ```
 

cmd/root.go

@@ -39,19 +39,20 @@ import (
 )
 
 var (
-	awsProfile  string
-	awsRegion   string
-	endpointURL string
-	local       bool
-	noColor     bool
-	mfa         bool
-	level       int
-	fullPath    bool
-	fileName    string
-	size 	  bool
+	awsProfile    string
+	awsRegion     string
+	endpointURL   string
+	local         bool
+	noColor       bool
+	mfa           bool
+	level         int
+	fullPath      bool
+	fileName      string
+	size          bool
 	humanReadable bool
-	dateTime   bool
-	username bool
+	dateTime      bool
+	username      bool
+	directoryOnly bool
 )
 
 var rootCmd = &cobra.Command{
@@ -82,18 +83,18 @@ var rootCmd = &cobra.Command{
 		if level > 0 {
 			maxDepth = &level
 		}
-		keys, err := pkg.FetchS3ObjectKeys(s3Svc, bucket, prefix, maxDepth,size, humanReadable,dateTime,username)
+		keys, err := pkg.FetchS3ObjectKeys(s3Svc, bucket, prefix, maxDepth, size, humanReadable, dateTime, username)
 		if err != nil {
 			log.Fatalf("failed to fetch S3 object keys: %v", err)
 			return
 		}
 
 		root := gtree.NewRoot(bucket)
 		if noColor || fileName != "" {
-			root = pkg.BuildTreeWithoutColor(root, bucket, keys, fullPath)
+			root = pkg.BuildTreeWithoutColor(root, bucket, keys, fullPath, directoryOnly)
 		} else {
 			root = gtree.NewRoot(color.BlueString(bucket))
-			root = pkg.BuildTreeWithColor(root, bucket, keys, fullPath)
+			root = pkg.BuildTreeWithColor(root, bucket, keys, fullPath, directoryOnly)
 		}
 
 		fileCount, dirCount := pkg.ProcessKeys(keys)
@@ -143,6 +144,7 @@ func init() {
 	rootCmd.Flags().BoolVarP(&humanReadable, "human-readable", "H", false, "Print the size of each file but in a more human readable way, e.g. appending a size letter for kilobytes (K), megabytes (M), gigabytes (G), terabytes (T), petabytes (P) and exabytes(E).")
 	rootCmd.Flags().BoolVarP(&dateTime, "date-time", "D", false, "Print the last modified time of each file.")
 	rootCmd.Flags().BoolVarP(&username, "username", "u", false, "Print the owner of each file.")
+	rootCmd.Flags().BoolVarP(&directoryOnly, "directory-only", "d", false, "List directories only.")
 }
 
 func extractBucketAndPrefix(input string) (string, string, error) {

pkg/s3.go

@@ -87,7 +87,7 @@ func formatBytes(b int64) string {
 }
 
 // FetchS3ObjectKeys returns a slice of keys for all objects in the specified bucket and prefix
-func FetchS3ObjectKeys(s3Client *s3.Client, bucket string, prefix string, maxDepth *int,size, humanReadable bool,dateTime bool,username bool) ([][]string, error) {
+func FetchS3ObjectKeys(s3Client *s3.Client, bucket string, prefix string, maxDepth *int, size, humanReadable bool, dateTime bool, username bool) ([][]string, error) {
 	var delimiter *string
 	var fetchOwner *bool
 	if maxDepth != nil {
@@ -115,9 +115,9 @@ func FetchS3ObjectKeys(s3Client *s3.Client, bucket string, prefix string, maxDep
 		}
 
 		input := &s3.ListObjectsV2Input{
-			Bucket:    aws.String(bucket),
-			Prefix:    aws.String(currentPrefix),
-			Delimiter: delimiter,
+			Bucket:     aws.String(bucket),
+			Prefix:     aws.String(currentPrefix),
+			Delimiter:  delimiter,
 			FetchOwner: fetchOwner,
 		}
 
@@ -131,8 +131,7 @@ func FetchS3ObjectKeys(s3Client *s3.Client, bucket string, prefix string, maxDep
 
 			for _, obj := range page.Contents {
 				key := strings.Split(*obj.Key, "/")
-				// fmt.Println(*obj.Owner.DisplayName)
-			
+
 				meta := []string{}
 
 				if username {

pkg/tree.go

@@ -18,29 +18,29 @@ func createFullPath(bucket string, keys [][]string) [][]string {
 }
 
 // BuildTreeWithColor builds a tree with colored nodes
-func BuildTreeWithColor(root *gtree.Node, bucket string, keys [][]string, f bool) *gtree.Node {
+func BuildTreeWithColor(root *gtree.Node, bucket string, keys [][]string, f bool, directoryOnly bool) *gtree.Node {
 	if f {
 		keys = createFullPath(bucket, keys)
 	}
 
 	for _, key := range keys {
-		addNodeWithColor(root, key, 0)
+		addNodeWithColor(root, key, 0, directoryOnly)
 	}
 	return root
 }
 
 // BuildTreeWithoutColor builds a tree without colored nodes
-func BuildTreeWithoutColor(root *gtree.Node, bucket string, keys [][]string, f bool) *gtree.Node {
+func BuildTreeWithoutColor(root *gtree.Node, bucket string, keys [][]string, f bool, directoryOnly bool) *gtree.Node {
 	if f {
 		keys = createFullPath(bucket, keys)
 	}
 	for _, key := range keys {
-		addNodeWithoutColor(root, key, 0)
+		addNodeWithoutColor(root, key, 0, directoryOnly)
 	}
 	return root
 }
 
-func addNodeWithoutColor(parent *gtree.Node, keys []string, depth int) *gtree.Node {
+func addNodeWithoutColor(parent *gtree.Node, keys []string, depth int, directoryOnly bool) *gtree.Node {
 	if len(keys) == 0 {
 		return nil
 	}
@@ -50,13 +50,18 @@ func addNodeWithoutColor(parent *gtree.Node, keys []string, depth int) *gtree.No
 		return nil
 	}
 
+	// Skip adding the root node if the tree is being built for a single directory
+	if directoryOnly && len(keys) == 1 {
+		return nil
+	}
+
 	// Add the current key (without color) as a node to the parent
 	node := parent.Add(keys[0])
 	// Recursively add the remaining keys
-	return addNodeWithoutColor(node, keys[1:], depth+1)
+	return addNodeWithoutColor(node, keys[1:], depth+1, directoryOnly)
 }
 
-func addNodeWithColor(parent *gtree.Node, keys []string, depth int) *gtree.Node {
+func addNodeWithColor(parent *gtree.Node, keys []string, depth int, directoryOnly bool) *gtree.Node {
 	if len(keys) == 0 {
 		return nil
 	}
@@ -66,11 +71,16 @@ func addNodeWithColor(parent *gtree.Node, keys []string, depth int) *gtree.Node
 		return nil
 	}
 
+	// Skip adding the root node if the tree is being built for a single directory
+	if directoryOnly && len(keys) == 1 {
+		return nil
+	}
+
 	// Add a colored node to the tree based on whether the current key represents a file or a directory
 	coloredKey := getColorizedKey(keys[0], len(keys) == 1)
 	node := parent.Add(coloredKey)
 
-	return addNodeWithColor(node, keys[1:], depth+1)
+	return addNodeWithColor(node, keys[1:], depth+1, directoryOnly)
 }
 
 // getColorizedKey returns the colored representation of a key, with different colors for files and directories