Add --max-depth flag to limit directory traversal depth

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
 

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,
 	)
 }

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,18 +80,19 @@ 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)
 		return
 	}
 
 	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,16 +117,19 @@ 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
 
 	dirColorFunc := getColorFunc(dirColorName)
 	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)
 			}

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)
 	}
 }
 

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.