codebase_doc.py

#!/usr/bin/env python3
import argparse
import re
from pathlib import Path
from typing import List, Dict, Pattern
from rich import print
from rich.console import Console
import pyperclip
from io import StringIO
from pathspec import PathSpec
from pathspec.patterns import GitWildMatchPattern

VERSION = "v0.0.1"


def create_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Generate markdown documentation of a codebase"
    )
    parser.add_argument("-i", "--input", default=".", help="Directory to scan")
    parser.add_argument(
        "-o", "--output", default="documentation.md", help="Output file name"
    )
    parser.add_argument(
        "-c",
        "--clipboard",
        action="store_true",
        help="Copy output to clipboard instead of writing to file",
    )
    parser.add_argument(
        "--ignore",
        default=r"\.git.*|__pycache__",
        help="Comma-separated list of regex patterns to ignore",
    )
    parser.add_argument(
        "--included-paths", help="File to save included paths (optional)"
    )
    parser.add_argument(
        "--excluded-paths", help="File to save excluded paths (optional)"
    )
    parser.add_argument(
        "--skip-gitignore", action="store_true", help="Skip loading .gitignore files"
    )
    parser.add_argument(
        "-v", "--version", action="store_true", help="Show version and exit"
    )
    return parser


def load_gitignore_patterns(base_path: Path) -> PathSpec:
    """Load all .gitignore files from the base path and its parents."""
    patterns = []

    # Start from base path and work up to root
    current_path = base_path.resolve()
    while True:
        gitignore_path = current_path / ".gitignore"
        if gitignore_path.is_file():
            try:
                with open(gitignore_path, "r", encoding="utf-8") as f:
                    # Add all non-empty, non-comment lines
                    patterns.extend(
                        line.strip()
                        for line in f
                        if line.strip() and not line.startswith("#")
                    )
            except Exception as e:
                print(
                    f"[yellow]Warning: Could not read {gitignore_path}: {str(e)}[/yellow]"
                )

        # Stop if we've reached the root
        if current_path.parent == current_path:
            break
        current_path = current_path.parent

    return PathSpec.from_lines(GitWildMatchPattern, patterns)


def should_ignore(
    path: str, ignore_patterns: List[Pattern], gitignore_spec: PathSpec
) -> bool:
    """Check if a path should be ignored based on regex patterns and gitignore rules."""
    if path == ".":
        return True

    # Check regex patterns first
    if any(pattern.search(str(path)) for pattern in ignore_patterns):
        return True

    # Then check gitignore patterns
    return gitignore_spec.match_file(str(path))


def get_relative_path(path: Path, base_path: Path) -> Path:
    try:
        return path.relative_to(base_path)
    except ValueError:
        return path


class OutputWriter:
    def __init__(self, file_path: Path = None, use_clipboard: bool = False):
        self.file_path = file_path
        self.use_clipboard = use_clipboard
        self.buffer = StringIO() if use_clipboard else None

        # If using file output, create/clear the file
        if file_path and not use_clipboard:
            with open(file_path, "w", encoding="utf-8"):
                pass

    def write(self, content: str) -> None:
        if self.use_clipboard:
            self.buffer.write(content)
        else:
            with open(self.file_path, "a", encoding="utf-8") as f:
                f.write(content)

    def finalize(self) -> None:
        if self.use_clipboard:
            pyperclip.copy(self.buffer.getvalue())
            self.buffer.close()


def print_tree(
    dir_path: Path,
    output: OutputWriter,
    ignore_patterns: List[Pattern],
    gitignore_spec: PathSpec,
    prefix: str = "",
    is_last: bool = True,
    depth_open: Dict[int, bool] = None,
) -> None:
    if depth_open is None:
        depth_open = {}

    try:
        entries = []
        for entry in sorted(dir_path.iterdir()):
            rel_path = get_relative_path(entry, Path.cwd())
            if not should_ignore(rel_path, ignore_patterns, gitignore_spec):
                entries.append(entry)

        for i, entry in enumerate(entries):
            is_last_entry = i == len(entries) - 1
            current_prefix = "└─" if is_last_entry else "├─"
            full_prefix = prefix + current_prefix
            entry_name = entry.name
            output.write(f"{full_prefix}{entry_name}\n")

            if entry.is_dir():
                next_prefix = prefix + ("  " if is_last_entry else "│ ")
                print_tree(
                    entry,
                    output,
                    ignore_patterns,
                    gitignore_spec,
                    next_prefix,
                    is_last_entry,
                )

    except PermissionError:
        print(f"[yellow]Warning: Permission denied accessing {dir_path}[/yellow]")


def write_code_content(
    dir_path: Path,
    output: OutputWriter,
    ignore_patterns: List[Pattern],
    gitignore_spec: PathSpec,
    included_paths_file: Path = None,
    excluded_paths_file: Path = None,
) -> None:
    console = Console()
    included_paths = []
    excluded_paths = []

    try:
        for path in sorted(dir_path.rglob("*")):
            rel_path = get_relative_path(path, Path.cwd())

            if should_ignore(rel_path, ignore_patterns, gitignore_spec):
                if excluded_paths_file:
                    excluded_paths.append(str(rel_path))
                else:
                    console.print(f"[red]- {rel_path}[/red]")
                continue

            if included_paths_file:
                included_paths.append(str(rel_path))
            else:
                console.print(f"[green]+ {rel_path}[/green]")

            if path.is_file():
                try:
                    content = path.read_text(errors="replace")
                    extension = path.suffix.lower().lstrip(".") or "txt"
                    output.write(f"\n## {rel_path}\n```{extension}\n{content}\n```\n\n")
                except Exception as e:
                    console.print(
                        f"[yellow]Warning: Could not read {rel_path}: {str(e)}[/yellow]"
                    )

        # Save path lists if requested
        if included_paths_file:
            with open(included_paths_file, "w", encoding="utf-8") as f:
                f.write("\n".join(included_paths))
        if excluded_paths_file:
            with open(excluded_paths_file, "w", encoding="utf-8") as f:
                f.write("\n".join(excluded_paths))

    except PermissionError:
        print(
            f"[yellow]Warning: Permission denied accessing some files in {dir_path}[/yellow]"
        )


def main():
    parser = create_parser()
    args = parser.parse_args()

    if args.version:
        print(VERSION)
        return

    # Convert ignore patterns to regex objects
    ignore_patterns = [
        re.compile(pattern.strip()) for pattern in args.ignore.split(",")
    ]

    # Print ignore patterns
    print("Regex patterns to ignore:")
    for pattern in ignore_patterns:
        print(f"  {pattern.pattern}")

    # Load gitignore patterns if not skipped
    input_path = Path(args.input).resolve()
    gitignore_spec = PathSpec.from_lines(
        GitWildMatchPattern, []
    )  # Empty spec by default

    if not args.skip_gitignore:
        gitignore_spec = load_gitignore_patterns(input_path)

        print("\nGitignore patterns loaded from:")
        current_path = input_path
        while True:
            gitignore_path = current_path / ".gitignore"
            if gitignore_path.is_file():
                print(f"  {gitignore_path}")
            if current_path.parent == current_path:
                break
            current_path = current_path.parent
    else:
        print("\nSkipping .gitignore files")

    # Setup output handling
    output_file = Path(args.output) if not args.clipboard else None
    output = OutputWriter(output_file, args.clipboard)

    # Write header
    output.write("# Tree View:\n```\n")

    # Generate tree structure
    print_tree(input_path, output, ignore_patterns, gitignore_spec)
    output.write("```\n\n# Content:\n")

    # Write file contents
    included_paths_file = Path(args.included_paths) if args.included_paths else None
    excluded_paths_file = Path(args.excluded_paths) if args.excluded_paths else None

    write_code_content(
        input_path,
        output,
        ignore_patterns,
        gitignore_spec,
        included_paths_file,
        excluded_paths_file,
    )

    # Finalize output (copies to clipboard if using clipboard mode)
    output.finalize()

    if args.clipboard:
        print("\nCodebase documentation copied to clipboard!")
    else:
        print("\nCodebase documentation generated successfully!")


if __name__ == "__main__":
    main()