From 589d0fb3f85c5ad2646d81cab20b78ae055a7966 Mon Sep 17 00:00:00 2001 From: zyadtaha Date: Sat, 1 Feb 2025 08:39:22 +0200 Subject: [PATCH] Add --max-depth flag to limit directory traversal depth --- README.md | 15 +++---- cmd/main.go | 10 +++++ pkg/printer/printer.go | 31 +++++++++------ pkg/printer/printer_bench_test.go | 6 +-- pkg/printer/printer_test.go | 66 +++++++++++++++++++++++++------ 5 files changed, 95 insertions(+), 33 deletions(-) diff --git a/README.md b/README.md index 8e9e555..d43567a 100644 --- a/README.md +++ b/README.md @@ -28,13 +28,14 @@ pr -dir /path/to/your/folder ### Basic Flags -| Flag | Description | Default | Example | -|------|-------------|---------|---------| -| `--dir` | Specify directory to print | Current directory | `pr --dir /path/to/folder` | -| `--ext` | Filter files by extension | All files | `pr --ext .go` | -| `--output` | Save output to file | Terminal output | `pr --output output.txt` | -| `--no-color` | Disable colored output | Colors enabled | `pr --no-color` | -| `--hidden` | Include hidden files | Not included | `pr --hidden` | +| Flag | Description | Default | Example | +|----------------|-------------|---------|-------| +| `--dir` | Specify directory to print | Current directory | `pr --dir /path/to/folder` | +| `--ext` | Filter files by extension | All files | `pr --ext .go` | +| `--output` | Save output to file | Terminal output | `pr --output output.txt` | +| `--no-color` | Disable colored output | Colors enabled | `pr --no-color` | +| `--hidden` | Include hidden files | Not included | `pr --hidden` | +| `--max-depth` | Limit directory traversal depth | No limit | `pr --max-depth 2` | ### Sorting Flags diff --git a/cmd/main.go b/cmd/main.go index 1d2e5c3..9660070 100644 --- a/cmd/main.go +++ b/cmd/main.go @@ -3,6 +3,8 @@ package main import ( "PrintLayout/pkg/printer" "flag" + "fmt" + "os" ) func main() { @@ -20,6 +22,7 @@ func main() { flag.StringVar(&config.SortBy, "sort-by", "name", "Sort by 'name', 'size', or 'time'") flag.StringVar(&config.Order, "order", "asc", "Sort order 'asc' or 'desc'") flag.BoolVar(&config.IncludeHidden, "hidden", false, "Include hidden files and directories") + flag.IntVar(&config.MaxDepth, "max-depth", -1, "Maximum depth of directory traversal") // Add --exclude flag to specify exclusion patterns flag.Func("exclude", "Exclude files/directories matching the pattern (can be specified multiple times)", func(pattern string) error { @@ -30,6 +33,12 @@ func main() { // Parse flags flag.Parse() + // Validate max-depth + if config.MaxDepth < -1 { + fmt.Fprintln(os.Stderr, "Error: --max-depth must be -1 (unlimited) or a non-negative integer.") + return + } + printer.PrintProjectStructure( config.DirPath, config.OutputPath, @@ -43,5 +52,6 @@ func main() { config.SortBy, config.Order, config.IncludeHidden, + config.MaxDepth, ) } diff --git a/pkg/printer/printer.go b/pkg/printer/printer.go index 475d038..9fa8417 100644 --- a/pkg/printer/printer.go +++ b/pkg/printer/printer.go @@ -27,6 +27,7 @@ type Config struct { SortBy string // "name", "size", "time" Order string // "asc", "desc" IncludeHidden bool + MaxDepth int } var colorMap = map[string]color.Attribute{ @@ -62,7 +63,8 @@ func HandleFlags(config Config) { config.ExcludePatterns, config.SortBy, config.Order, - config.IncludeHidden) + config.IncludeHidden, + config.MaxDepth) } // PrintProjectStructure prints the directory structure of the given root directory. @@ -78,7 +80,8 @@ func PrintProjectStructure( excludePatterns []string, sortBy string, order string, - includeHidden bool) { + includeHidden bool, + maxDepth int) { absRoot, err := filepath.Abs(root) if err != nil { fmt.Println("Error getting absolute path:", err) @@ -86,10 +89,10 @@ func PrintProjectStructure( } if format == "text" { - dirCount, fileCount := getTreeOutput(absRoot, extFilter, useColor, dirColorName, fileColorName, execColorName, excludePatterns, sortBy, order, includeHidden) + dirCount, fileCount := getTreeOutput(absRoot, extFilter, useColor, dirColorName, fileColorName, execColorName, excludePatterns, sortBy, order, includeHidden, maxDepth) fmt.Printf("\n%d directories, %d files\n", dirCount, fileCount) } else { - tree := buildTree(absRoot, extFilter, excludePatterns, sortBy, order, includeHidden) + tree := buildTree(absRoot, extFilter, excludePatterns, sortBy, order, includeHidden, maxDepth, 0) var output string switch format { case "json": @@ -114,7 +117,7 @@ func PrintProjectStructure( } } -func getTreeOutput(root string, extFilter string, useColor bool, dirColorName string, fileColorName string, execColorName string, excludePatterns []string, sortBy string, order string, includeHidden bool) (int, int) { +func getTreeOutput(root string, extFilter string, useColor bool, dirColorName string, fileColorName string, execColorName string, excludePatterns []string, sortBy string, order string, includeHidden bool, maxDepth int) (int, int) { dirCount := 0 fileCount := 0 @@ -122,8 +125,11 @@ func getTreeOutput(root string, extFilter string, useColor bool, dirColorName st fileColorFunc := getColorFunc(fileColorName) execColorFunc := getColorFunc(execColorName) - var traverse func(string, string) error - traverse = func(currentDir string, prefix string) error { + var traverse func(string, string, int) error + traverse = func(currentDir string, prefix string, depth int) error { + if maxDepth != -1 && depth >= maxDepth { + return nil + } dir, err := os.Open(currentDir) if err != nil { return err @@ -160,7 +166,7 @@ func getTreeOutput(root string, extFilter string, useColor bool, dirColorName st fmt.Printf("%s%s/\n", prefix+getTreePrefix(isLast), entry.Name()) } - err := traverse(filepath.Join(currentDir, entry.Name()), prefix+getIndent(isLast)) + err := traverse(filepath.Join(currentDir, entry.Name()), prefix+getIndent(isLast), depth+1) if err != nil { return err } @@ -187,7 +193,7 @@ func getTreeOutput(root string, extFilter string, useColor bool, dirColorName st } fmt.Printf("%s/\n", filepath.Base(root)) - err := traverse(root, "") + err := traverse(root, "", 0) if err != nil { fmt.Println("Error traversing directory:", err) } @@ -243,7 +249,10 @@ type Node struct { } // buildTree constructs a tree of Nodes from the directory structure -func buildTree(currentDir string, extFilter string, excludePatterns []string, sortBy string, order string, includeHidden bool) *Node { +func buildTree(currentDir string, extFilter string, excludePatterns []string, sortBy string, order string, includeHidden bool, maxDepth int, depth int) *Node { + if maxDepth != -1 && depth >= maxDepth { + return nil + } dir, err := os.Open(currentDir) if err != nil { return nil @@ -274,7 +283,7 @@ func buildTree(currentDir string, extFilter string, excludePatterns []string, so } if entry.IsDir() { - child := buildTree(filepath.Join(currentDir, entry.Name()), extFilter, excludePatterns, sortBy, order, includeHidden) + child := buildTree(filepath.Join(currentDir, entry.Name()), extFilter, excludePatterns, sortBy, order, includeHidden, maxDepth, depth+1) if child != nil { node.Children = append(node.Children, child) } diff --git a/pkg/printer/printer_bench_test.go b/pkg/printer/printer_bench_test.go index 537a252..2f3db8a 100644 --- a/pkg/printer/printer_bench_test.go +++ b/pkg/printer/printer_bench_test.go @@ -24,7 +24,7 @@ func BenchmarkPrintProjectStructure(b *testing.B) { // Run the benchmark b.ResetTimer() // Reset the timer to exclude setup time for i := 0; i < b.N; i++ { - PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", false) + PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", false, -1) } } @@ -43,7 +43,7 @@ func BenchmarkPrintProjectStructure_JSON(b *testing.B) { b.ResetTimer() for i := 0; i < b.N; i++ { - PrintProjectStructure(".", "", "", false, "json", "blue", "green", "red", []string{}, "name", "asc", false) + PrintProjectStructure(".", "", "", false, "json", "blue", "green", "red", []string{}, "name", "asc", false, -1) } } @@ -61,7 +61,7 @@ func BenchmarkPrintProjectStructure_LargeDirectory(b *testing.B) { b.ResetTimer() // Reset the timer to exclude setup time for i := 0; i < b.N; i++ { - PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", false) + PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", false, -1) } } diff --git a/pkg/printer/printer_test.go b/pkg/printer/printer_test.go index 21f28ea..5c93ba6 100644 --- a/pkg/printer/printer_test.go +++ b/pkg/printer/printer_test.go @@ -29,7 +29,7 @@ func TestPrintProjectStructure(t *testing.T) { // Test text output t.Run("TextOutput", func(t *testing.T) { output := captureOutput(func() { - PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", false) + PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", false, -1) }) rootName := filepath.Base(tmpDir) @@ -57,7 +57,7 @@ func TestPrintProjectStructure(t *testing.T) { // Test JSON output t.Run("JSONOutput", func(t *testing.T) { output := captureOutput(func() { - PrintProjectStructure(".", "", "", false, "json", "blue", "green", "red", []string{}, "name", "asc", false) + PrintProjectStructure(".", "", "", false, "json", "blue", "green", "red", []string{}, "name", "asc", false, -1) }) // Verify that the output is valid JSON @@ -70,7 +70,7 @@ func TestPrintProjectStructure(t *testing.T) { // Test XML output t.Run("XMLOutput", func(t *testing.T) { output := captureOutput(func() { - PrintProjectStructure(".", "", "", false, "xml", "blue", "green", "red", []string{}, "name", "asc", false) + PrintProjectStructure(".", "", "", false, "xml", "blue", "green", "red", []string{}, "name", "asc", false, -1) }) // Verify that the output is valid XML @@ -83,7 +83,7 @@ func TestPrintProjectStructure(t *testing.T) { // Test YAML output t.Run("YAMLOutput", func(t *testing.T) { output := captureOutput(func() { - PrintProjectStructure(".", "", "", false, "yaml", "blue", "green", "red", []string{}, "name", "asc", false) + PrintProjectStructure(".", "", "", false, "yaml", "blue", "green", "red", []string{}, "name", "asc", false, -1) }) // Verify that the output is valid YAML @@ -96,7 +96,7 @@ func TestPrintProjectStructure(t *testing.T) { // Test exclusion patterns t.Run("ExclusionPatterns", func(t *testing.T) { output := captureOutput(func() { - PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{"*.go"}, "name", "asc", false) + PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{"*.go"}, "name", "asc", false, -1) }) rootName := filepath.Base(tmpDir) @@ -120,7 +120,7 @@ func TestPrintProjectStructure(t *testing.T) { // Test sorting by name (ascending) t.Run("SortByNameAsc", func(t *testing.T) { output := captureOutput(func() { - PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", false) + PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", false, -1) }) // Verify that the output is sorted by name in ascending order @@ -131,7 +131,7 @@ func TestPrintProjectStructure(t *testing.T) { // Test sorting by name (descending) t.Run("SortByNameDesc", func(t *testing.T) { output := captureOutput(func() { - PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "desc", false) + PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "desc", false, -1) }) // Verify that the output is sorted by name in descending order @@ -142,7 +142,7 @@ func TestPrintProjectStructure(t *testing.T) { // Test sorting by size (ascending) t.Run("SortBySizeAsc", func(t *testing.T) { output := captureOutput(func() { - PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "size", "asc", false) + PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "size", "asc", false, -1) }) // Verify that the output is sorted by size in ascending order @@ -153,7 +153,7 @@ func TestPrintProjectStructure(t *testing.T) { // Test sorting by size (descending) t.Run("SortBySizeDesc", func(t *testing.T) { output := captureOutput(func() { - PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "size", "desc", false) + PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "size", "desc", false, -1) }) // Verify that the output is sorted by size in descending order @@ -164,7 +164,7 @@ func TestPrintProjectStructure(t *testing.T) { // Test sorting by time (ascending) t.Run("SortByTimeAsc", func(t *testing.T) { output := captureOutput(func() { - PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "time", "asc", false) + PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "time", "asc", false, -1) }) // Verify that the output is sorted by time in ascending order @@ -175,7 +175,7 @@ func TestPrintProjectStructure(t *testing.T) { // Test sorting by time (descending) t.Run("SortByTimeDesc", func(t *testing.T) { output := captureOutput(func() { - PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "time", "desc", false) + PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "time", "desc", false, -1) }) // Verify that the output is sorted by time in descending order @@ -186,13 +186,55 @@ func TestPrintProjectStructure(t *testing.T) { // Test including hidden files t.Run("IncludeHidden", func(t *testing.T) { output := captureOutput(func() { - PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", true) + PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", true, -1) }) // Verify that the output includes the hidden files // You can add specific checks based on your expected output t.Log(output) }) + + // Test maximum depth 0 (should only print the root directory) + t.Run("MaxDepthZero", func(t *testing.T) { + output := captureOutput(func() { + PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", false, 0) + }) + + rootName := filepath.Base(tmpDir) + expected := rootName + "/\n" + + "\n0 directories, 0 files\n" + output = strings.TrimSpace(output) + expected = strings.TrimSpace(expected) + + if output != expected { + t.Errorf("Unexpected output:\nGot:\n%s\nExpected:\n%s", output, expected) + } + }) + + // Test maximum depth 2 + t.Run("MaxDepthTwo", func(t *testing.T) { + output := captureOutput(func() { + PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", false, 2) + }) + + rootName := filepath.Base(tmpDir) + + expected := rootName + "/\n" + + "├── cmd/\n" + + "│ └── main.go\n" + + "├── go.mod\n" + + "├── internal/\n" + + "│ └── utils/\n" + + "└── pkg/\n" + + " └── printer/\n" + + "\n5 directories, 2 files\n" + output = strings.TrimSpace(output) + expected = strings.TrimSpace(expected) + + if output != expected { + t.Errorf("Unexpected output:\nGot:\n%s\nExpected:\n%s", output, expected) + } + }) } // createTestProjectStructure creates a sample project structure for testing.