Skip to content

Commit 743ca7b

Browse files
AhmedhossamdevAhmedhossamdev
authored
Merge pull request #21 from zyadtaha/max-depth-flag
Add --max-depth flag to limit directory traversal depth
2 parents d1e7be8 + 589d0fb commit 743ca7b

File tree

5 files changed

+95
-33
lines changed

5 files changed

+95
-33
lines changed

‎README.md‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -28,13 +28,14 @@ pr -dir /path/to/your/folder
2828

2929
### Basic Flags
3030

31-
| Flag | Description | Default | Example |
32-
|------|-------------|---------|---------|
33-
| `--dir` | Specify directory to print | Current directory | `pr --dir /path/to/folder` |
34-
| `--ext` | Filter files by extension | All files | `pr --ext .go` |
35-
| `--output` | Save output to file | Terminal output | `pr --output output.txt` |
36-
| `--no-color` | Disable colored output | Colors enabled | `pr --no-color` |
37-
| `--hidden` | Include hidden files | Not included | `pr --hidden` |
31+
| Flag | Description | Default | Example |
32+
|----------------|-------------|---------|-------|
33+
| `--dir` | Specify directory to print | Current directory | `pr --dir /path/to/folder` |
34+
| `--ext` | Filter files by extension | All files | `pr --ext .go` |
35+
| `--output` | Save output to file | Terminal output | `pr --output output.txt` |
36+
| `--no-color` | Disable colored output | Colors enabled | `pr --no-color` |
37+
| `--hidden` | Include hidden files | Not included | `pr --hidden` |
38+
| `--max-depth` | Limit directory traversal depth | No limit | `pr --max-depth 2` |
3839

3940
### Sorting Flags
4041

‎cmd/main.go‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@ package main
33
import (
44
"PrintLayout/pkg/printer"
55
"flag"
6+
"fmt"
7+
"os"
68
)
79

810
func main() {
@@ -20,6 +22,7 @@ func main() {
2022
flag.StringVar(&config.SortBy, "sort-by", "name", "Sort by 'name', 'size', or 'time'")
2123
flag.StringVar(&config.Order, "order", "asc", "Sort order 'asc' or 'desc'")
2224
flag.BoolVar(&config.IncludeHidden, "hidden", false, "Include hidden files and directories")
25+
flag.IntVar(&config.MaxDepth, "max-depth", -1, "Maximum depth of directory traversal")
2326

2427
// Add --exclude flag to specify exclusion patterns
2528
flag.Func("exclude", "Exclude files/directories matching the pattern (can be specified multiple times)", func(pattern string) error {
@@ -30,6 +33,12 @@ func main() {
3033
// Parse flags
3134
flag.Parse()
3235

36+
// Validate max-depth
37+
if config.MaxDepth < -1 {
38+
fmt.Fprintln(os.Stderr, "Error: --max-depth must be -1 (unlimited) or a non-negative integer.")
39+
return
40+
}
41+
3342
printer.PrintProjectStructure(
3443
config.DirPath,
3544
config.OutputPath,
@@ -43,5 +52,6 @@ func main() {
4352
config.SortBy,
4453
config.Order,
4554
config.IncludeHidden,
55+
config.MaxDepth,
4656
)
4757
}

‎pkg/printer/printer.go‎

Lines changed: 20 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ type Config struct {
2727
SortBy string // "name", "size", "time"
2828
Order string // "asc", "desc"
2929
IncludeHidden bool
30+
MaxDepth int
3031
}
3132

3233
var colorMap = map[string]color.Attribute{
@@ -62,7 +63,8 @@ func HandleFlags(config Config) {
6263
config.ExcludePatterns,
6364
config.SortBy,
6465
config.Order,
65-
config.IncludeHidden)
66+
config.IncludeHidden,
67+
config.MaxDepth)
6668
}
6769

6870
// PrintProjectStructure prints the directory structure of the given root directory.
@@ -78,18 +80,19 @@ func PrintProjectStructure(
7880
excludePatterns []string,
7981
sortBy string,
8082
order string,
81-
includeHidden bool) {
83+
includeHidden bool,
84+
maxDepth int) {
8285
absRoot, err := filepath.Abs(root)
8386
if err != nil {
8487
fmt.Println("Error getting absolute path:", err)
8588
return
8689
}
8790

8891
if format == "text" {
89-
dirCount, fileCount := getTreeOutput(absRoot, extFilter, useColor, dirColorName, fileColorName, execColorName, excludePatterns, sortBy, order, includeHidden)
92+
dirCount, fileCount := getTreeOutput(absRoot, extFilter, useColor, dirColorName, fileColorName, execColorName, excludePatterns, sortBy, order, includeHidden, maxDepth)
9093
fmt.Printf("\n%d directories, %d files\n", dirCount, fileCount)
9194
} else {
92-
tree := buildTree(absRoot, extFilter, excludePatterns, sortBy, order, includeHidden)
95+
tree := buildTree(absRoot, extFilter, excludePatterns, sortBy, order, includeHidden, maxDepth, 0)
9396
var output string
9497
switch format {
9598
case "json":
@@ -114,16 +117,19 @@ func PrintProjectStructure(
114117
}
115118
}
116119

117-
func getTreeOutput(root string, extFilter string, useColor bool, dirColorName string, fileColorName string, execColorName string, excludePatterns []string, sortBy string, order string, includeHidden bool) (int, int) {
120+
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) {
118121
dirCount := 0
119122
fileCount := 0
120123

121124
dirColorFunc := getColorFunc(dirColorName)
122125
fileColorFunc := getColorFunc(fileColorName)
123126
execColorFunc := getColorFunc(execColorName)
124127

125-
var traverse func(string, string) error
126-
traverse = func(currentDir string, prefix string) error {
128+
var traverse func(string, string, int) error
129+
traverse = func(currentDir string, prefix string, depth int) error {
130+
if maxDepth != -1 && depth >= maxDepth {
131+
return nil
132+
}
127133
dir, err := os.Open(currentDir)
128134
if err != nil {
129135
return err
@@ -160,7 +166,7 @@ func getTreeOutput(root string, extFilter string, useColor bool, dirColorName st
160166
fmt.Printf("%s%s/\n", prefix+getTreePrefix(isLast), entry.Name())
161167
}
162168

163-
err := traverse(filepath.Join(currentDir, entry.Name()), prefix+getIndent(isLast))
169+
err := traverse(filepath.Join(currentDir, entry.Name()), prefix+getIndent(isLast), depth+1)
164170
if err != nil {
165171
return err
166172
}
@@ -187,7 +193,7 @@ func getTreeOutput(root string, extFilter string, useColor bool, dirColorName st
187193
}
188194

189195
fmt.Printf("%s/\n", filepath.Base(root))
190-
err := traverse(root, "")
196+
err := traverse(root, "", 0)
191197
if err != nil {
192198
fmt.Println("Error traversing directory:", err)
193199
}
@@ -243,7 +249,10 @@ type Node struct {
243249
}
244250

245251
// buildTree constructs a tree of Nodes from the directory structure
246-
func buildTree(currentDir string, extFilter string, excludePatterns []string, sortBy string, order string, includeHidden bool) *Node {
252+
func buildTree(currentDir string, extFilter string, excludePatterns []string, sortBy string, order string, includeHidden bool, maxDepth int, depth int) *Node {
253+
if maxDepth != -1 && depth >= maxDepth {
254+
return nil
255+
}
247256
dir, err := os.Open(currentDir)
248257
if err != nil {
249258
return nil
@@ -274,7 +283,7 @@ func buildTree(currentDir string, extFilter string, excludePatterns []string, so
274283
}
275284

276285
if entry.IsDir() {
277-
child := buildTree(filepath.Join(currentDir, entry.Name()), extFilter, excludePatterns, sortBy, order, includeHidden)
286+
child := buildTree(filepath.Join(currentDir, entry.Name()), extFilter, excludePatterns, sortBy, order, includeHidden, maxDepth, depth+1)
278287
if child != nil {
279288
node.Children = append(node.Children, child)
280289
}

‎pkg/printer/printer_bench_test.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ func BenchmarkPrintProjectStructure(b *testing.B) {
2424
// Run the benchmark
2525
b.ResetTimer() // Reset the timer to exclude setup time
2626
for i := 0; i < b.N; i++ {
27-
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", false)
27+
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", false, -1)
2828
}
2929
}
3030

@@ -43,7 +43,7 @@ func BenchmarkPrintProjectStructure_JSON(b *testing.B) {
4343

4444
b.ResetTimer()
4545
for i := 0; i < b.N; i++ {
46-
PrintProjectStructure(".", "", "", false, "json", "blue", "green", "red", []string{}, "name", "asc", false)
46+
PrintProjectStructure(".", "", "", false, "json", "blue", "green", "red", []string{}, "name", "asc", false, -1)
4747
}
4848
}
4949

@@ -61,7 +61,7 @@ func BenchmarkPrintProjectStructure_LargeDirectory(b *testing.B) {
6161

6262
b.ResetTimer() // Reset the timer to exclude setup time
6363
for i := 0; i < b.N; i++ {
64-
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", false)
64+
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", false, -1)
6565
}
6666
}
6767

‎pkg/printer/printer_test.go‎

Lines changed: 54 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ func TestPrintProjectStructure(t *testing.T) {
2929
// Test text output
3030
t.Run("TextOutput", func(t *testing.T) {
3131
output := captureOutput(func() {
32-
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", false)
32+
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", false, -1)
3333
})
3434

3535
rootName := filepath.Base(tmpDir)
@@ -57,7 +57,7 @@ func TestPrintProjectStructure(t *testing.T) {
5757
// Test JSON output
5858
t.Run("JSONOutput", func(t *testing.T) {
5959
output := captureOutput(func() {
60-
PrintProjectStructure(".", "", "", false, "json", "blue", "green", "red", []string{}, "name", "asc", false)
60+
PrintProjectStructure(".", "", "", false, "json", "blue", "green", "red", []string{}, "name", "asc", false, -1)
6161
})
6262

6363
// Verify that the output is valid JSON
@@ -70,7 +70,7 @@ func TestPrintProjectStructure(t *testing.T) {
7070
// Test XML output
7171
t.Run("XMLOutput", func(t *testing.T) {
7272
output := captureOutput(func() {
73-
PrintProjectStructure(".", "", "", false, "xml", "blue", "green", "red", []string{}, "name", "asc", false)
73+
PrintProjectStructure(".", "", "", false, "xml", "blue", "green", "red", []string{}, "name", "asc", false, -1)
7474
})
7575

7676
// Verify that the output is valid XML
@@ -83,7 +83,7 @@ func TestPrintProjectStructure(t *testing.T) {
8383
// Test YAML output
8484
t.Run("YAMLOutput", func(t *testing.T) {
8585
output := captureOutput(func() {
86-
PrintProjectStructure(".", "", "", false, "yaml", "blue", "green", "red", []string{}, "name", "asc", false)
86+
PrintProjectStructure(".", "", "", false, "yaml", "blue", "green", "red", []string{}, "name", "asc", false, -1)
8787
})
8888

8989
// Verify that the output is valid YAML
@@ -96,7 +96,7 @@ func TestPrintProjectStructure(t *testing.T) {
9696
// Test exclusion patterns
9797
t.Run("ExclusionPatterns", func(t *testing.T) {
9898
output := captureOutput(func() {
99-
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{"*.go"}, "name", "asc", false)
99+
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{"*.go"}, "name", "asc", false, -1)
100100
})
101101

102102
rootName := filepath.Base(tmpDir)
@@ -120,7 +120,7 @@ func TestPrintProjectStructure(t *testing.T) {
120120
// Test sorting by name (ascending)
121121
t.Run("SortByNameAsc", func(t *testing.T) {
122122
output := captureOutput(func() {
123-
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", false)
123+
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", false, -1)
124124
})
125125

126126
// Verify that the output is sorted by name in ascending order
@@ -131,7 +131,7 @@ func TestPrintProjectStructure(t *testing.T) {
131131
// Test sorting by name (descending)
132132
t.Run("SortByNameDesc", func(t *testing.T) {
133133
output := captureOutput(func() {
134-
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "desc", false)
134+
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "desc", false, -1)
135135
})
136136

137137
// Verify that the output is sorted by name in descending order
@@ -142,7 +142,7 @@ func TestPrintProjectStructure(t *testing.T) {
142142
// Test sorting by size (ascending)
143143
t.Run("SortBySizeAsc", func(t *testing.T) {
144144
output := captureOutput(func() {
145-
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "size", "asc", false)
145+
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "size", "asc", false, -1)
146146
})
147147

148148
// Verify that the output is sorted by size in ascending order
@@ -153,7 +153,7 @@ func TestPrintProjectStructure(t *testing.T) {
153153
// Test sorting by size (descending)
154154
t.Run("SortBySizeDesc", func(t *testing.T) {
155155
output := captureOutput(func() {
156-
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "size", "desc", false)
156+
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "size", "desc", false, -1)
157157
})
158158

159159
// Verify that the output is sorted by size in descending order
@@ -164,7 +164,7 @@ func TestPrintProjectStructure(t *testing.T) {
164164
// Test sorting by time (ascending)
165165
t.Run("SortByTimeAsc", func(t *testing.T) {
166166
output := captureOutput(func() {
167-
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "time", "asc", false)
167+
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "time", "asc", false, -1)
168168
})
169169

170170
// Verify that the output is sorted by time in ascending order
@@ -175,7 +175,7 @@ func TestPrintProjectStructure(t *testing.T) {
175175
// Test sorting by time (descending)
176176
t.Run("SortByTimeDesc", func(t *testing.T) {
177177
output := captureOutput(func() {
178-
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "time", "desc", false)
178+
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "time", "desc", false, -1)
179179
})
180180

181181
// Verify that the output is sorted by time in descending order
@@ -186,13 +186,55 @@ func TestPrintProjectStructure(t *testing.T) {
186186
// Test including hidden files
187187
t.Run("IncludeHidden", func(t *testing.T) {
188188
output := captureOutput(func() {
189-
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", true)
189+
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", true, -1)
190190
})
191191

192192
// Verify that the output includes the hidden files
193193
// You can add specific checks based on your expected output
194194
t.Log(output)
195195
})
196+
197+
// Test maximum depth 0 (should only print the root directory)
198+
t.Run("MaxDepthZero", func(t *testing.T) {
199+
output := captureOutput(func() {
200+
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", false, 0)
201+
})
202+
203+
rootName := filepath.Base(tmpDir)
204+
expected := rootName + "/\n" +
205+
"\n0 directories, 0 files\n"
206+
output = strings.TrimSpace(output)
207+
expected = strings.TrimSpace(expected)
208+
209+
if output != expected {
210+
t.Errorf("Unexpected output:\nGot:\n%s\nExpected:\n%s", output, expected)
211+
}
212+
})
213+
214+
// Test maximum depth 2
215+
t.Run("MaxDepthTwo", func(t *testing.T) {
216+
output := captureOutput(func() {
217+
PrintProjectStructure(".", "", "", false, "text", "blue", "green", "red", []string{}, "name", "asc", false, 2)
218+
})
219+
220+
rootName := filepath.Base(tmpDir)
221+
222+
expected := rootName + "/\n" +
223+
"├── cmd/\n" +
224+
"│ └── main.go\n" +
225+
"├── go.mod\n" +
226+
"├── internal/\n" +
227+
"│ └── utils/\n" +
228+
"└── pkg/\n" +
229+
" └── printer/\n" +
230+
"\n5 directories, 2 files\n"
231+
output = strings.TrimSpace(output)
232+
expected = strings.TrimSpace(expected)
233+
234+
if output != expected {
235+
t.Errorf("Unexpected output:\nGot:\n%s\nExpected:\n%s", output, expected)
236+
}
237+
})
196238
}
197239

198240
// createTestProjectStructure creates a sample project structure for testing.

0 commit comments

Comments
 (0)