add omit test files and only frontend

README.md

@@ -13,8 +13,8 @@ Generate markdown documentation from directory structure. Optimized for LLM inge
 - Force overwrite of existing output files
 - Configurable depth, file size limits, ignore patterns
 - Optional upload to Proton Drive after generating output
-
+- Skip frontend files (`.html`, `.css`, `.js`, etc.)
-![dirmd](./docs/dirmd.gif)
+- Omit test files and inline test blocks (Go, Java, Rust)
 
 ## Build
 
@@ -37,12 +37,49 @@ Generate markdown documentation from directory structure. Optimized for LLM inge
     # Generate with vertical slices
     dirmd -o output.md ~/repos/myproject --vertical-slices
 
+    # Skip frontend files (.html, .css, .js, etc.)
+    dirmd -o output.md ~/repos/myproject --skip-frontend
+
+    # Omit test files and inline test blocks
+    # Go: excludes *_test.go
+    # Java: excludes *Test.java, *Tests.java, src/test/java/*
+    # Rust: strips #[cfg(test)] mod tests { ... } blocks
+    dirmd -o output.md ~/repos/myproject --omit-test
+
+    # Combine filters
+    dirmd -o output.md ~/repos/myproject --skip-frontend --omit-test
+
     # Generate and upload to Proton Drive (default: /my-files/md/<filename>)
     dirmd -o output.md ~/repos/myproject --proton-drive
 
     # Generate and upload to a specific remote path
     dirmd -o output.md ~/repos/myproject --proton-drive --drive-path /my-files/docs/api.md
 
+## Test File Omission
+
+The `--omit-test` flag removes test files and inline test blocks to reduce token overhead in LLM context windows. Behavior varies by language:
+
+| Language | Exclusion Strategy | What Gets Removed |
+|----------|-------------------|-------------------|
+| **Go** | File exclusion | All `*_test.go` files |
+| **Java** | File exclusion | `*Test.java`, `*Tests.java`, files under `src/test/java/` |
+| **Rust** | Content stripping | Entire `#[cfg(test)] mod tests { ... }` blocks removed from `.rs` files |
+
+The Rust stripping works by detecting the `#[cfg(test)]` attribute, then tracking brace depth until the module closes. Nested braces in string literals may occasionally throw off the count, but this is acceptable for RAG use cases.
+
+Frontend files can be excluded separately with `--skip-frontend`:
+
+| Extension Type | Excluded With `--skip-frontend` |
+|---------------|-------------------------------|
+| Templates | `.html`, `.htm`, `.gohtml`, `.tmpl` |
+| Stylesheets | `.css`, `.scss`, `.sass`, `.less` |
+| Scripts | `.js`, `.jsx`, `.mjs`, `.cjs`, `.ts`, `.tsx` |
+| Frameworks | `.vue`, `.svelte` |
+
+These flags can be combined freely:
+
+    dirmd -o docs.md ~/repos/myproject --skip-frontend --omit-test
+
 ## Vertical Slices Mode
 
 When `--vertical-slices` is passed, dirmd splits output into multiple files instead of one monolithic document:
@@ -84,7 +121,9 @@ Output is designed for RAG ingestion and LLM context windows, not human presenta
 | `--ignore` | Additional ignore patterns (glob) | |
 | `--extensions` | Additional file extensions to skip | |
 | `--proton-drive` | Upload output to Proton Drive after writing locally | false |
-| `--drive-path` | Full remote path on Proton Drive (e.g. /my-files/md/report.md) | /my-files/md/&lt;filename&gt; |
+| `--drive-path` | Full remote path on Proton Drive (e.g. /my-files/md/report.md) | /my-files/md/<filename> |
+| `--skip-frontend` | Omit frontend file types (.html, .css, .js, .ts, .vue, .svelte, etc.) | false |
+| `--omit-test` | Exclude test files (Go *_test.go, Java *Test.java) and strip #[cfg(test)] blocks from Rust | false |
 
 ## Proton Drive Integration
 
@@ -96,7 +135,7 @@ In vertical slices mode, all generated files (root + slices) are uploaded to `/m
 
 1. Install the Proton Drive CLI — see the [official guide](https://proton.me/support/drive-cli) for download and installation instructions.
 2. Authenticate by running `proton-drive auth login`.
-3. Ensure `proton-drive` is in your `PATH`.
+3. Ensure `proton-drive` is in your PATH.
 
 If the CLI is not installed or you are not logged in, dirmd will report the error to stderr and exit with a non-zero code. Local files are always written successfully before the upload is attempted.
 
@@ -116,6 +155,9 @@ For full documentation on the Proton Drive CLI, including installation, authenti
     # Vertical slices with overwrite
     dirmd -o docs.md ~/repos/project -f --vertical-slices
 
+    # Skip frontend and omit tests
+    dirmd -o docs.md ~/repos/project -f --skip-frontend --omit-test
+
     # Generate, then upload to default remote location
     dirmd -o docs.md ~/repos/project --proton-drive
 

cmd/root.go

@@ -1,3 +1,4 @@
+// cmd/root.go
 package cmd
 
 import (
@@ -41,6 +42,9 @@ Examples:
   # Skip frontend files
   dirmd -o ./docs ~/repos/myproject --skip-frontend
 
+  # Omit test files and test blocks
+  dirmd -o ./docs ~/repos/myproject --omit-test
+
   # Generate and upload to Proton Drive
   dirmd -o ./docs ~/repos/myproject --proton-drive
 
@@ -61,6 +65,7 @@ Examples:
 	drivePath      string
 	verticalSlices bool
 	skipFrontend   bool
+	omitTest       bool
 )
 
 func Execute() {
@@ -82,6 +87,7 @@ func init() {
 	rootCmd.Flags().StringVar(&drivePath, "drive-path", "", "full remote directory on Proton Drive (e.g. /my-files/docs/api). Requires --proton-drive")
 	rootCmd.Flags().BoolVar(&verticalSlices, "vertical-slices", false, "split output into root index.md + one markdown per top-level directory")
 	rootCmd.Flags().BoolVar(&skipFrontend, "skip-frontend", false, "omit frontend file types (.html, .css, .js, .ts, .vue, .svelte, etc.)")
+	rootCmd.Flags().BoolVar(&omitTest, "omit-test", false, "exclude test files (Go *_test.go, Java *Test.java) and strip #[cfg(test)] blocks from Rust")
 
 	rootCmd.MarkPersistentFlagRequired("output")
 }
@@ -116,6 +122,7 @@ func run(cmd *cobra.Command, args []string) error {
 		DrivePath:      drivePath,
 		VerticalSlices: verticalSlices,
 		SkipFrontend:   skipFrontend,
+		OmitTest:       omitTest,
 	}
 
 	if singleFile != "" {

internal/config/config.go

@@ -1,3 +1,4 @@
+// internal/config/config.go
 package config
 
 import (
@@ -79,6 +80,7 @@ type Config struct {
 	DrivePath      string
 	VerticalSlices bool
 	SkipFrontend   bool
+	OmitTest       bool
 }
 
 func PrintUsage() {
@@ -104,6 +106,7 @@ func Parse() (*Config, error) {
 		drivePath      string
 		verticalSlices bool
 		skipFrontend   bool
+		omitTest       bool
 	)
 
 	flag.StringVar(&input, "i", "", "single file input (append mode only)")
@@ -119,6 +122,7 @@ func Parse() (*Config, error) {
 	flag.StringVar(&drivePath, "drive-path", "", "full remote path on Proton Drive (requires --proton-drive)")
 	flag.BoolVar(&verticalSlices, "vertical-slices", false, "split output into root file + one markdown per top-level directory")
 	flag.BoolVar(&skipFrontend, "skip-frontend", false, "omit frontend file types (.html, .css, .js, .ts, .vue, .svelte, etc.)")
+	flag.BoolVar(&omitTest, "omit-test", false, "exclude test files (Go *_test.go, Java *Test.java) and strip #[cfg(test)] blocks from Rust")
 
 	flag.Usage = PrintUsage
 	flag.Parse()
@@ -161,6 +165,7 @@ func Parse() (*Config, error) {
 		DrivePath:      drivePath,
 		VerticalSlices: verticalSlices,
 		SkipFrontend:   skipFrontend,
+		OmitTest:       omitTest,
 	}
 
 	if input != "" {

internal/filter/filter.go

@@ -1,3 +1,4 @@
+// internal/filter/filter.go
 package filter
 
 import (
@@ -10,6 +11,7 @@ type Filter struct {
 	MaxSize    int64
 	Extensions []string
 	Ignores    []string
+	OmitTest   bool
 }
 
 func (f *Filter) ShouldSkipDir(name string) bool {
@@ -48,6 +50,28 @@ func (f *Filter) ShouldSkipFile(path, name string, info os.FileInfo) bool {
 		return true
 	}
 
+	if f.OmitTest && isTestFile(path, name) {
+		return true
+	}
+
+	return false
+}
+
+func isTestFile(path, name string) bool {
+	ext := strings.ToLower(filepath.Ext(path))
+
+	switch ext {
+	case ".go":
+		return strings.HasSuffix(strings.ToLower(name), "_test.go")
+	case ".java":
+		base := strings.TrimSuffix(name, filepath.Ext(name))
+		if strings.HasSuffix(base, "Test") || strings.HasSuffix(base, "Tests") {
+			return true
+		}
+		slashPath := filepath.ToSlash(path)
+		return strings.Contains(slashPath, "/src/test/")
+	}
+
 	return false
 }
 

internal/renderer/renderer.go

@@ -1,3 +1,4 @@
+// internal/renderer/renderer.go
 package renderer
 
 import (
@@ -14,7 +15,6 @@ import (
 
 const fenceChars = "~~~"
 
-// Write writes a single markdown file with tree and file contents.
 func Write(result *walker.Result, outputPath string, appendMode bool, singleFile bool) error {
 	flags := os.O_CREATE | os.O_WRONLY
 	if appendMode {
@@ -55,7 +55,7 @@ func Write(result *walker.Result, outputPath string, appendMode bool, singleFile
 		for _, entry := range result.Entries {
 			fullPath := filepath.Join(result.AbsRoot, entry.RelPath)
 
-			content, err := readFile(fullPath)
+			content, err := readFileForOutput(fullPath, result.OmitTest)
 			if err != nil {
 				fmt.Fprintf(os.Stderr, "warning: cannot read %s: %v\n", entry.RelPath, err)
 				continue
@@ -82,8 +82,6 @@ func Write(result *walker.Result, outputPath string, appendMode bool, singleFile
 	return err
 }
 
-// WriteSlices generates a root markdown file plus one markdown per top-level directory.
-// Returns the list of generated file paths.
 func WriteSlices(result *walker.Result, outputDir, repoName string) ([]string, error) {
 	entries := result.Entries
 	if len(entries) == 0 {
@@ -109,16 +107,14 @@ func WriteSlices(result *walker.Result, outputDir, repoName string) ([]string, e
 
 	var generated []string
 
-	// Write root file as index_<reponame>.md
 	indexBasename := "index_" + repoName + ".md"
 	rootPath := filepath.Join(outputDir, indexBasename)
 	generated = append(generated, rootPath)
 
-	if err := writeRootSlice(rootPath, absRoot, rootName, repoName, rootFiles, topLevelDirs); err != nil {
+	if err := writeRootSlice(rootPath, absRoot, rootName, repoName, rootFiles, topLevelDirs, result.OmitTest); err != nil {
 		return nil, fmt.Errorf("failed to write root slice: %w", err)
 	}
 
-	// Write each top-level directory slice into a mirroring subdirectory
 	var dirNames []string
 	for k := range topLevelDirs {
 		dirNames = append(dirNames, k)
@@ -137,7 +133,6 @@ func WriteSlices(result *walker.Result, outputDir, repoName string) ([]string, e
 		absDir := filepath.Join(absRoot, dirName)
 		dirEntries := topLevelDirs[dirName]
 
-		// Strip dirName prefix for tree building only
 		treeEntries := make([]walker.Entry, len(dirEntries))
 		for i, e := range dirEntries {
 			sepIdx := strings.Index(e.RelPath, string(filepath.Separator))
@@ -151,7 +146,7 @@ func WriteSlices(result *walker.Result, outputDir, repoName string) ([]string, e
 			}
 		}
 
-		if err := writeDirSlice(slicePath, absDir, absRoot, dirName, treeEntries, dirEntries, indexBasename); err != nil {
+		if err := writeDirSlice(slicePath, absDir, absRoot, dirName, treeEntries, dirEntries, indexBasename, result.OmitTest); err != nil {
 			return nil, fmt.Errorf("failed to write slice %s: %w", sliceName, err)
 		}
 	}
@@ -159,7 +154,7 @@ func WriteSlices(result *walker.Result, outputDir, repoName string) ([]string, e
 	return generated, nil
 }
 
-func writeRootSlice(path, absRoot, rootName, repoName string, rootFiles []walker.Entry, topLevelDirs map[string][]walker.Entry) error {
+func writeRootSlice(path, absRoot, rootName, repoName string, rootFiles []walker.Entry, topLevelDirs map[string][]walker.Entry, omitTest bool) error {
 	f, err := os.Create(path)
 	if err != nil {
 		return fmt.Errorf("cannot create root file: %w", err)
@@ -170,14 +165,12 @@ func writeRootSlice(path, absRoot, rootName, repoName string, rootFiles []walker
 
 	fmt.Fprintf(&sb, "# %s [ROOT]\n\n", absRoot)
 
-	// Build combined tree with references
 	treeLines := buildTreeWithRefs(rootName, repoName, topLevelDirs, rootFiles)
 	sb.WriteString("## Directory Tree\n\n")
 	sb.WriteString(fenceChars + "\n")
 	sb.WriteString(treeLines)
 	sb.WriteString(fenceChars + "\n\n")
 
-	// Vertical slices index for LLM navigation
 	sb.WriteString("## Vertical Slices Index\n\n")
 	sb.WriteString("These files contain the full contents of each top-level directory.\n")
 	sb.WriteString("Read the relevant slice file for detailed source code.\n\n")
@@ -211,7 +204,7 @@ func writeRootSlice(path, absRoot, rootName, repoName string, rootFiles []walker
 
 		for _, entry := range rootFiles {
 			fullPath := filepath.Join(absRoot, entry.RelPath)
-			content, err := readFile(fullPath)
+			content, err := readFileForOutput(fullPath, omitTest)
 			if err != nil {
 				fmt.Fprintf(os.Stderr, "warning: cannot read %s: %v\n", entry.RelPath, err)
 				continue
@@ -238,7 +231,7 @@ func writeRootSlice(path, absRoot, rootName, repoName string, rootFiles []walker
 	return err
 }
 
-func writeDirSlice(path, absDir, absRoot, dirName string, treeEntries, fileEntries []walker.Entry, rootFileName string) error {
+func writeDirSlice(path, absDir, absRoot, dirName string, treeEntries, fileEntries []walker.Entry, rootFileName string, omitTest bool) error {
 	f, err := os.Create(path)
 	if err != nil {
 		return fmt.Errorf("cannot create slice file: %w", err)
@@ -266,7 +259,7 @@ func writeDirSlice(path, absDir, absRoot, dirName string, treeEntries, fileEntri
 
 		for _, entry := range fileEntries {
 			fullPath := filepath.Join(absRoot, entry.RelPath)
-			content, err := readFile(fullPath)
+			content, err := readFileForOutput(fullPath, omitTest)
 			if err != nil {
 				fmt.Fprintf(os.Stderr, "warning: cannot read %s: %v\n", entry.RelPath, err)
 				continue
@@ -298,19 +291,16 @@ func buildTreeWithRefs(rootName, repoName string, topLevelDirs map[string][]walk
 	sb.WriteString(rootName)
 	sb.WriteString("/\n")
 
-	// Collect and sort directory names
 	var dirNames []string
 	for k := range topLevelDirs {
 		dirNames = append(dirNames, k)
 	}
 	sort.Strings(dirNames)
 
-	// Sort root files
 	sort.Slice(rootFiles, func(i, j int) bool {
 		return rootFiles[i].RelPath < rootFiles[j].RelPath
 	})
 
-	// Merge directories and files into a single sorted list for display
 	type treeItem struct {
 		name     string
 		isDir    bool
@@ -441,6 +431,68 @@ func readFile(path string) (string, error) {
 	return content, nil
 }
 
+func readFileForOutput(fullPath string, omitTest bool) (string, error) {
+	content, err := readFile(fullPath)
+	if err != nil {
+		return "", err
+	}
+	if omitTest && strings.EqualFold(filepath.Ext(fullPath), ".rs") {
+		content = stripTestBlocks(content)
+	}
+	return content, nil
+}
+
+func stripTestBlocks(content string) string {
+	lines := strings.Split(content, "\n")
+	var result []string
+
+	i := 0
+	for i < len(lines) {
+		line := lines[i]
+		trimmed := strings.TrimSpace(line)
+
+		if isCfgTestAttr(trimmed) {
+			braceDepth := 0
+			foundBrace := false
+
+			for i < len(lines) {
+				line := lines[i]
+
+				if !foundBrace {
+					if strings.Contains(line, "{") {
+						braceDepth += strings.Count(line, "{") - strings.Count(line, "}")
+						foundBrace = true
+					}
+					i++
+					continue
+				}
+
+				braceDepth += strings.Count(line, "{") - strings.Count(line, "}")
+				i++
+
+				if braceDepth <= 0 {
+					break
+				}
+			}
+			continue
+		}
+
+		result = append(result, line)
+		i++
+	}
+
+	if len(result) == 0 {
+		return ""
+	}
+	return strings.Join(result, "\n")
+}
+
+func isCfgTestAttr(line string) bool {
+	return strings.Contains(line, "#[cfg(test)]") ||
+		strings.Contains(line, "#[cfg(all(test") ||
+		strings.Contains(line, "#[cfg(any(test")
+}
+
 func extToLang(ext string) string {
 	mapping := map[string]string{
 		".go":   "go",

internal/walker/walker.go

@@ -1,3 +1,4 @@
+// internal/walker/walker.go
 package walker
 
 import (
@@ -20,6 +21,7 @@ type Result struct {
 	Entries    []Entry
 	AbsRoot    string
 	SingleFile bool
+	OmitTest   bool
 }
 
 func Run(cfg *config.Config) (*Result, error) {
@@ -43,6 +45,7 @@ func processSingleFile(cfg *config.Config) (*Result, error) {
 		MaxSize:    cfg.MaxSize,
 		Extensions: cfg.Extensions,
 		Ignores:    cfg.Ignores,
+		OmitTest:   cfg.OmitTest,
 	}
 
 	if flt.ShouldSkipFile(cfg.AbsRoot, filepath.Base(cfg.AbsRoot), info) {
@@ -53,8 +56,8 @@ func processSingleFile(cfg *config.Config) (*Result, error) {
 		Entries:    []Entry{{RelPath: filepath.Base(cfg.AbsRoot), Size: info.Size()}},
 		AbsRoot:    filepath.Dir(cfg.AbsRoot),
 		SingleFile: true,
+		OmitTest:   cfg.OmitTest,
 	}, nil
-
 }
 
 func walkDirectory(cfg *config.Config) (*Result, error) {
@@ -62,9 +65,10 @@ func walkDirectory(cfg *config.Config) (*Result, error) {
 		MaxSize:    cfg.MaxSize,
 		Extensions: cfg.Extensions,
 		Ignores:    cfg.Ignores,
+		OmitTest:   cfg.OmitTest,
 	}
 
-	result := &Result{AbsRoot: cfg.AbsRoot}
+	result := &Result{AbsRoot: cfg.AbsRoot, OmitTest: cfg.OmitTest}
 
 	err := filepath.WalkDir(cfg.AbsRoot, func(path string, d fs.DirEntry, err error) error {
 		if err != nil {