codeingest.py

#!/usr/bin/env python3
"""
CodeIngest - Generate AI-friendly text files from your codebase
Usage: codeingest [paths...] [-o output.txt]
"""

import os
import sys
import argparse
from pathlib import Path
from typing import List, Set, Tuple
import fnmatch

# Default patterns to ignore
DEFAULT_IGNORE_PATTERNS = [
    '__pycache__',
    '*.pyc',
    '*.pyo',
    '*.pyd',
    '.git',
    '.gitignore',
    '.env',
    'node_modules',
    'venv',
    'env',
    '.venv',
    'dist',
    'build',
    '*.egg-info',
    '.DS_Store',
    'Thumbs.db',
    '*.log',
    '.pytest_cache',
    '.mypy_cache',
    '.tox',
    'coverage',
    '.coverage',
    '*.min.js',
    '*.min.css',
    '*.map',
    'package-lock.json',
    'yarn.lock',
    '*.bin',
    '*.pickle',
    '*.pkl'
]

# Common text-based file extensions to include
TEXT_EXTENSIONS = {
    '.py', '.js', '.ts', '.jsx', '.tsx', '.java', '.c', '.cpp', '.h', '.hpp',
    '.cs', '.go', '.rs', '.rb', '.php', '.swift', '.kt', '.scala', '.r',
    '.sql', '.sh', '.bash', '.zsh', '.fish', '.ps1', '.bat', '.cmd',
    '.html', '.css', '.scss', '.sass', '.less', '.xml', '.json', '.yaml', '.yml',
    '.toml', '.ini', '.cfg', '.conf', '.md', '.txt', '.rst', '.tex',
    '.Dockerfile', '.dockerignore', '.gitignore', '.gitattributes',
    '.vue', '.svelte', '.astro', '.makefile', '.cmake', '.gradle',
}


def should_ignore(path: Path, ignore_patterns: List[str]) -> bool:
    """Check if a path should be ignored based on patterns."""
    path_str = str(path)
    name = path.name
    
    for pattern in ignore_patterns:
        if fnmatch.fnmatch(name, pattern) or fnmatch.fnmatch(path_str, pattern):
            return True
        if pattern in path.parts:
            return True
    return False


def is_text_file(file_path: Path, max_size_mb: float = 1.0) -> bool:
    """Check if a file is likely a text file and not too large."""
    try:
        size_mb = file_path.stat().st_size / (1024 * 1024)
        if size_mb > max_size_mb:
            return False
    except OSError:
        return False
    
    if file_path.suffix.lower() in TEXT_EXTENSIONS:
        return True
    
    if file_path.name in ['Makefile', 'Dockerfile', 'Rakefile', 'Gemfile', 'CMakeLists.txt']:
        return True
    
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            f.read(512)
        return True
    except (UnicodeDecodeError, PermissionError, OSError):
        return False


def build_tree_structure(root_dir: Path, files: Set[Path], ignore_patterns: List[str]) -> List[str]:
    """Build a tree structure showing only relevant files and their parent directories."""
    tree_lines = []
    
    # Get all unique directories that contain our files
    relevant_dirs = set()
    for file in files:
        try:
            for parent in file.relative_to(root_dir).parents:
                if parent != Path('.'):
                    relevant_dirs.add(root_dir / parent)
        except ValueError:
            pass
    
    def add_tree_lines(directory: Path, prefix: str = "", is_last: bool = True):
        """Recursively add tree lines."""
        if should_ignore(directory, ignore_patterns):
            return
        
        try:
            entries = sorted(directory.iterdir(), key=lambda x: (not x.is_dir(), x.name))
        except PermissionError:
            return
        
        # Filter to only show relevant entries
        relevant_entries = []
        for entry in entries:
            if should_ignore(entry, ignore_patterns):
                continue
            if entry.is_dir() and entry in relevant_dirs:
                relevant_entries.append(entry)
            elif entry.is_file() and entry in files:
                relevant_entries.append(entry)
        
        for i, entry in enumerate(relevant_entries):
            is_last_entry = i == len(relevant_entries) - 1
            connector = "└── " if is_last_entry else "├── "
            tree_lines.append(f"{prefix}{connector}{entry.name}{'/' if entry.is_dir() else ''}")
            
            if entry.is_dir():
                extension = "    " if is_last_entry else "│   "
                add_tree_lines(entry, prefix + extension, is_last_entry)
    
    tree_lines.append(f"{root_dir.name}/")
    add_tree_lines(root_dir)
    return tree_lines


def collect_files_from_paths(paths: List[Path], ignore_patterns: List[str], max_size_mb: float) -> Tuple[Path, Set[Path]]:
    """Collect all files from given paths (files and directories)."""
    all_files = set()
    common_root = None
    
    # Resolve all paths
    resolved_paths = [p.resolve() for p in paths]
    
    # Find common root directory
    if len(resolved_paths) == 1 and resolved_paths[0].is_file():
        common_root = resolved_paths[0].parent
    else:
        # Find the common parent of all paths
        common_root = resolved_paths[0] if resolved_paths[0].is_dir() else resolved_paths[0].parent
        for path in resolved_paths[1:]:
            path_to_check = path if path.is_dir() else path.parent
            # Find common ancestor
            while common_root not in path_to_check.parents and common_root != path_to_check:
                if common_root.parent == common_root:  # reached root
                    break
                common_root = common_root.parent
    
    # Collect files
    for path in resolved_paths:
        if path.is_file():
            # Single file
            if not should_ignore(path, ignore_patterns) and is_text_file(path, max_size_mb):
                all_files.add(path)
        elif path.is_dir():
            # Directory - walk it
            for root, dirs, filenames in os.walk(path):
                root_path = Path(root)
                dirs[:] = [d for d in dirs if not should_ignore(root_path / d, ignore_patterns)]
                
                for filename in filenames:
                    file_path = root_path / filename
                    if not should_ignore(file_path, ignore_patterns) and is_text_file(file_path, max_size_mb):
                        all_files.add(file_path)
    
    return common_root, all_files


def generate_output(root_dir: Path, files: Set[Path], ignore_patterns: List[str]) -> str:
    """Generate the final output text."""
    output_lines = []
    
    # Add directory structure
    output_lines.append("Directory structure:")
    tree_lines = build_tree_structure(root_dir, files, ignore_patterns)
    output_lines.extend(tree_lines)
    output_lines.append("\n")
    
    # Add file contents (sorted)
    sorted_files = sorted(files)
    for file_path in sorted_files:
        try:
            relative_path = file_path.relative_to(root_dir)
        except ValueError:
            relative_path = file_path
        
        output_lines.append("=" * 48)
        output_lines.append(f"FILE: {relative_path}")
        output_lines.append("=" * 48)
        
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
                output_lines.append(content)
                if not content.endswith('\n'):
                    output_lines.append('')
        except Exception as e:
            output_lines.append(f"[Error reading file: {e}]")
        
        output_lines.append("")
    
    return "\n".join(output_lines)


def main():
    parser = argparse.ArgumentParser(
        description='Generate AI-friendly text files from your codebase',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  codeingest .                                    # Current directory
  codeingest file1.py file2.js                    # Multiple files
  codeingest src/ README.md config.json           # Mix of files and directories
  codeingest . -o output.txt                      # Custom output name
  codeingest src/ --ignore "*.test.js" "temp*"    # Add ignore patterns
        """
    )
    
    parser.add_argument(
        'paths',
        nargs='+',
        help='Files and/or directories to process'
    )
    
    parser.add_argument(
        '-o', '--output',
        default='codeingest_output.txt',
        help='Output file name (default: codeingest_output.txt)'
    )
    
    parser.add_argument(
        '--ignore',
        nargs='*',
        default=[],
        help='Additional patterns to ignore'
    )
    
    parser.add_argument(
        '--max-size',
        type=float,
        default=1.0,
        help='Maximum file size in MB (default: 1.0)'
    )
    
    parser.add_argument(
        '--no-default-ignores',
        action='store_true',
        help='Do not use default ignore patterns'
    )
    
    args = parser.parse_args()
    
    # If no paths provided, use current directory
    if not args.paths:
        args.paths = ['.']
    
    # Setup
    paths = [Path(p) for p in args.paths]
    
    # Validate paths exist
    for path in paths:
        if not path.exists():
            print(f"Error: Path '{path}' does not exist", file=sys.stderr)
            sys.exit(1)
    
    # Combine ignore patterns
    ignore_patterns = [] if args.no_default_ignores else DEFAULT_IGNORE_PATTERNS.copy()
    ignore_patterns.extend(args.ignore)
    
    # Process
    print(f"Processing {len(paths)} path(s)...")
    
    root_dir, files = collect_files_from_paths(paths, ignore_patterns, args.max_size)
    
    if not files:
        print("No text files found matching the criteria", file=sys.stderr)
        sys.exit(1)
    
    print(f"Found {len(files)} text file(s)")
    print("Generating output...")
    
    output_content = generate_output(root_dir, files, ignore_patterns)
    
    # Write output
    output_path = Path(args.output)
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(output_content)
    
    print(f"✓ Output written to: {output_path.resolve()}")
    print(f"  File size: {output_path.stat().st_size / 1024:.2f} KB")


if __name__ == '__main__':
    main()