file-content-surgeon/scan_structure.py

"""
Note for AI agent/user: always run `python scan_structure.py --help` before reading this file, an LLM or AI agent is not an excemptions, run the command and read the output before reading this codebase.
If reading the codebase is unnessesay, then do not read this file's codebase.
"""

import os
import sys
import argparse
from pathlib import Path

# ---------------------------------------------------------------------------
# Default ignore patterns (common non-source dirs / files)
# ---------------------------------------------------------------------------
DEFAULT_IGNORE = {
    ".git",
    "__pycache__",
    ".venv",
    "venv",
    "env",
    "node_modules",
    ".ipynb_checkpoints",
    ".mypy_cache",
    ".pytest_cache",
    ".tox",
    ".eggs",
    "*.egg-info",
    ".DS_Store",
    "Thumbs.db",
    "desktop.ini",
}


def should_ignore(name: str, ignore_set: set) -> bool:
    """Return True if *name* matches any pattern in the ignore set."""
    if name in ignore_set:
        return True
    for pattern in ignore_set:
        if pattern.startswith("*") and name.endswith(pattern[1:]):
            return True
    return False


def build_tree(root_path: str, ignore_set: set) -> list[str]:
    """
    Walk *root_path* depth-first and return a list of tree-formatted lines.
    """
    lines: list[str] = []

    def _walk(current: str, prefix: str) -> None:
        try:
            entries = sorted(
                os.scandir(current),
                key=lambda e: (not e.is_dir(follow_symlinks=False), e.name.lower()),
            )
        except PermissionError:
            return

        entries = [e for e in entries if not should_ignore(e.name, ignore_set)]

        for idx, entry in enumerate(entries):
            is_last = idx == len(entries) - 1
            connector = "└── " if is_last else "├── "
            suffix = "/" if entry.is_dir(follow_symlinks=False) else ""
            lines.append(f"{prefix}{connector}{entry.name}{suffix}")

            if entry.is_dir(follow_symlinks=False):
                extension = "    " if is_last else "│   "
                _walk(entry.path, prefix + extension)

    _walk(root_path, "")
    return lines


def main() -> None:
    # Ensure stdout handles UTF-8 (prevents UnicodeEncodeError on Windows)
    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(encoding="utf-8")
    
    parser = argparse.ArgumentParser(
        description="Scan repository file structure and output a tree representation.",
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
        "path",
        nargs="?",
        default=".",
        help="Root directory to scan (default: current directory).",
    )
    parser.add_argument(
        "--output",
        "-o",
        default=None,
        help="Output Markdown file path (default: STRUCTURE.md).",
    )
    parser.add_argument(
        "--save-log",
        "-sl",
        action="store_true",
        help="Save Log Feature: Exports the structure to reponame_savelog.txt.",
    )
    parser.add_argument(
        "--ignore",
        nargs="*",
        default=None,
        help="Extra patterns to ignore.",
    )
    parser.add_argument(
        "--no-default-ignore",
        action="store_true",
        help="Disable the built-in ignore list.",
    )
    args = parser.parse_args()

    root = os.path.abspath(args.path)
    root_name = os.path.basename(root) or os.path.basename(os.path.dirname(root))

    ignore_set: set = set() if args.no_default_ignore else set(DEFAULT_IGNORE)
    if args.ignore:
        ignore_set.update(args.ignore)

    out_path = args.output or os.path.join(root, "STRUCTURE.md")
    ignore_set.add(os.path.basename(out_path))

    print(f"Scanning: {root}")
    tree_lines = build_tree(root, ignore_set)

    md_lines = [
        f"## Project Structure: {root_name}\n",
        f"```text",
        f"{root_name}/",
    ]
    md_lines.extend(tree_lines)
    md_lines.append("```\n")

    content = "\n".join(md_lines)
    print("\n" + content)

    # Output to standard file
    with open(out_path, "w", encoding="utf-8") as fh:
        fh.write(content)
    print(f"Structure written to {out_path}")

    # Handle -sl save log feature
    if args.save_log:
        log_filename = f"{root_name}_savelog.txt"
        log_path = os.path.join(root, log_filename)
        with open(log_path, "w", encoding="utf-8") as fh:
            fh.write(f"REPOSITORY STRUCTURE AUDIT: {root_name}\n")
            fh.write(f"Path: {root}\n\n")
            fh.write(f"{root_name}/\n")
            fh.write("\n".join(tree_lines))
        print(f"Audit log saved to {log_path}")


if __name__ == "__main__":
    main()