feat: add whitelist file support (-w/--whitelist-file, .treemapperwhitelist)

Author: nikolay-e

src/treemapper/__init__.py

@@ -5,7 +5,7 @@
 from typing import Any
 
 from .diffctx import build_diff_context
-from .ignore import get_ignore_specs
+from .ignore import get_ignore_specs, get_whitelist_spec
 from .tree import TreeBuildContext, build_tree
 from .version import __version__
 from .writer import write_tree_json, write_tree_markdown, write_tree_text, write_tree_yaml
@@ -31,12 +31,14 @@ def map_directory(
     max_file_bytes: int | None = None,
     ignore_file: str | Path | None = None,
     no_default_ignores: bool = False,
+    whitelist_file: str | Path | None = None,
 ) -> dict[str, Any]:
     root_dir = Path(path).resolve()
     if not root_dir.is_dir():
         raise ValueError(f"'{path}' is not a directory")
 
     ignore_path = Path(ignore_file).resolve() if ignore_file else None
+    whitelist_path = Path(whitelist_file).resolve() if whitelist_file else None
 
     ctx = TreeBuildContext(
         base_dir=root_dir,
@@ -45,6 +47,7 @@ def map_directory(
         max_depth=max_depth,
         no_content=no_content,
         max_file_bytes=max_file_bytes,
+        whitelist_spec=get_whitelist_spec(whitelist_path, root_dir),
     )
 
     return {

src/treemapper/cli.py

@@ -81,10 +81,20 @@ def _resolve_ignore_file(ignore_file_arg: str | None) -> Path | None:
     return ignore_file
 
 
+def _resolve_whitelist_file(whitelist_file_arg: str | None) -> Path | None:
+    if not whitelist_file_arg:
+        return None
+    whitelist_file = Path(whitelist_file_arg).resolve()
+    if not whitelist_file.is_file():
+        _exit_error(f"Whitelist file '{whitelist_file_arg}' does not exist.")
+    return whitelist_file
+
+
 @dataclass
 class ParsedArgs:
     root_dir: Path
     ignore_file: Path | None
+    whitelist_file: Path | None
     output_file: Path | None
     no_default_ignores: bool
     verbosity: int
@@ -132,6 +142,7 @@ def parse_args() -> ParsedArgs:
     parser.add_argument("-v", "--version", action="version", version=f"%(prog)s {__version__}")
     parser.add_argument("directory", nargs="?", default=".", help="The directory to analyze")
     parser.add_argument("-i", "--ignore-file", default=None, help="Path to custom ignore file")
+    parser.add_argument("-w", "--whitelist-file", default=None, help="Path to whitelist file (only matching files are included)")
     parser.add_argument(
         "-o",
         "--output-file",
@@ -211,13 +222,15 @@ def parse_args() -> ParsedArgs:
     output_format = "yaml" if args.format == "yml" else args.format
     output_file, force_stdout = _resolve_output_file(args.output_file, output_format)
     ignore_file = _resolve_ignore_file(args.ignore_file)
+    whitelist_file = _resolve_whitelist_file(args.whitelist_file)
 
     log_level_map = {"error": 0, "warning": 1, "info": 2, "debug": 3}
     verbosity = log_level_map[args.log_level]
 
     return ParsedArgs(
         root_dir=root_dir,
         ignore_file=ignore_file,
+        whitelist_file=whitelist_file,
         output_file=output_file,
         no_default_ignores=args.no_default_ignores,
         verbosity=verbosity,

src/treemapper/diffctx/__init__.py

@@ -8,7 +8,7 @@
 
 import pathspec
 
-from ..ignore import get_ignore_specs, should_ignore
+from ..ignore import get_ignore_specs, get_whitelist_spec, is_whitelisted, should_ignore
 from ..tokens import count_tokens
 from .config import LIMITS
 from .config.extensions import CODE_EXTENSIONS, CONFIG_EXTENSIONS, DOC_EXTENSIONS
@@ -220,6 +220,7 @@ def build_diff_context(
     ignore_file: Path | None = None,
     no_default_ignores: bool = False,
     full: bool = False,
+    whitelist_file: Path | None = None,
 ) -> dict[str, Any]:
     _validate_inputs(root_dir, alpha, tau, budget_tokens)
     root_dir = root_dir.resolve()
@@ -229,6 +230,7 @@ def build_diff_context(
     base_rev, head_rev = split_diff_range(diff_range)
     is_working_tree_diff = base_rev is None and head_rev is None
     combined_spec = get_ignore_specs(root_dir, ignore_file, no_default_ignores, None)
+    wl_spec = get_whitelist_spec(whitelist_file, root_dir)
 
     untracked: list[Path] = []
     if is_working_tree_diff:
@@ -247,13 +249,15 @@ def build_diff_context(
     changed_files = [_normalize_path(p, root_dir) for p in changed_files]
     changed_files.extend(untracked)
     changed_files = _filter_ignored(changed_files, root_dir, combined_spec)
+    changed_files = _filter_whitelist(changed_files, root_dir, wl_spec)
 
     preferred_revs = _build_preferred_revs(base_rev, head_rev)
 
     seen_frag_ids: set[FragmentId] = set()
     all_fragments = _process_files_for_fragments(changed_files, root_dir, preferred_revs, seen_frag_ids)
 
     all_candidate_files = _collect_candidate_files(root_dir, set(changed_files), combined_spec)
+    all_candidate_files = _filter_whitelist(all_candidate_files, root_dir, wl_spec)
 
     edge_discovered = discover_all_related_files(changed_files, all_candidate_files, root_dir)
     edge_discovered = [_normalize_path(p, root_dir) for p in edge_discovered]
@@ -586,6 +590,24 @@ def _collect_expansion_files(
     return list(expansion_files)
 
 
+def _filter_whitelist(
+    files: list[Path],
+    root_dir: Path,
+    wl_spec: pathspec.PathSpec | None,
+) -> list[Path]:
+    if wl_spec is None:
+        return files
+    result: list[Path] = []
+    for file_path in files:
+        try:
+            rel_path = file_path.relative_to(root_dir).as_posix()
+            if is_whitelisted(rel_path, wl_spec):
+                result.append(file_path)
+        except ValueError:
+            pass
+    return result
+
+
 def _filter_ignored(
     files: list[Path],
     root_dir: Path,

src/treemapper/ignore.py

@@ -272,3 +272,28 @@ def should_ignore(relative_path_str: str, combined_spec: pathspec.PathSpec) -> b
     if logging.getLogger().isEnabledFor(logging.DEBUG):
         logging.debug("Checking ignore for '%s': %s", relative_path_str, is_ignored)
     return is_ignored
+
+
+DEFAULT_WHITELIST_FILENAME = ".treemapperwhitelist"
+
+
+def get_whitelist_spec(whitelist_file: Path | None, root_dir: Path | None = None) -> pathspec.PathSpec | None:
+    effective_file = whitelist_file
+    if not effective_file and root_dir:
+        default = root_dir / DEFAULT_WHITELIST_FILENAME
+        if default.is_file():
+            effective_file = default
+    if not effective_file:
+        return None
+    patterns = read_ignore_file(effective_file)
+    if not patterns:
+        return None
+    return pathspec.PathSpec.from_lines("gitignore", patterns)
+
+
+def is_whitelisted(relative_path_str: str, whitelist_spec: pathspec.PathSpec | None, is_dir: bool = False) -> bool:
+    if whitelist_spec is None:
+        return True
+    if is_dir:
+        return True
+    return whitelist_spec.match_file(relative_path_str)

src/treemapper/tree.py

@@ -7,7 +7,7 @@
 
 import pathspec
 
-from .ignore import should_ignore
+from .ignore import is_whitelisted, should_ignore
 
 BINARY_DETECTION_SAMPLE_SIZE = 8192
 MAX_SAFE_FILE_SIZE = 100 * 1024 * 1024  # 100 MB - prevent OOM when --max-file-bytes 0
@@ -108,6 +108,7 @@ class TreeBuildContext:
     max_depth: int | None = None
     no_content: bool = False
     max_file_bytes: int | None = None
+    whitelist_spec: pathspec.PathSpec | None = None
     _resolved_output_file: Path | None = None
 
     def __post_init__(self) -> None:
@@ -161,6 +162,9 @@ def _process_entry(entry: Path, ctx: TreeBuildContext, current_depth: int) -> di
     if should_ignore(path_to_check, ctx.combined_spec):
         return None
 
+    if not is_whitelisted(relative_path, ctx.whitelist_spec, is_dir=is_dir):
+        return None
+
     if entry.is_symlink() or not entry.exists():
         logging.debug("Skipping '%s': symlink or not exists", path_to_check)
         return None
@@ -176,6 +180,8 @@ def _create_node(entry: Path, ctx: TreeBuildContext, current_depth: int, is_dir:
             children = build_tree(entry, ctx, current_depth + 1)
             if children:
                 node["children"] = children
+            elif ctx.whitelist_spec is not None:
+                return None
         elif not ctx.no_content:
             node["content"] = _read_file_content(entry, ctx.max_file_bytes)
 

src/treemapper/treemapper.py

@@ -28,14 +28,15 @@ def _build_diff_tree(args: ParsedArgs) -> dict[str, Any]:
             ignore_file=args.ignore_file,
             no_default_ignores=args.no_default_ignores,
             full=args.full_diff,
+            whitelist_file=args.whitelist_file,
         )
     except GitError as e:
         print(f"Error: {e}", file=sys.stderr)
         sys.exit(1)
 
 
 def _build_standard_tree(args: ParsedArgs) -> dict[str, Any]:
-    from .ignore import get_ignore_specs
+    from .ignore import get_ignore_specs, get_whitelist_spec
     from .tree import TreeBuildContext, build_tree
 
     ctx = TreeBuildContext(
@@ -45,6 +46,7 @@ def _build_standard_tree(args: ParsedArgs) -> dict[str, Any]:
         max_depth=args.max_depth,
         no_content=args.no_content,
         max_file_bytes=args.max_file_bytes,
+        whitelist_spec=get_whitelist_spec(args.whitelist_file, args.root_dir),
     )
     return {
         "name": args.root_dir.name,

src/treemapper/writer.py

@@ -75,7 +75,7 @@ def _write_yaml_content(file: TextIO, content: str, base_indent: str) -> None:
     content_indent = base_indent + "  "
     if not content:
         file.write(f'{base_indent}content: ""\n')
-    elif _has_problematic_chars(content):
+    elif _has_problematic_chars(content) or not content.strip():
         file.write(f'{base_indent}content: "{_escape_yaml_content(content)}"\n')
     else:
         file.write(f"{base_indent}content: |2\n")