flatten_directory.py

#!/usr/bin/env python3
"""
Flatten Directory - A tool to flatten all files in a directory into a single file.
Inspired by rendergit (https://github.com/karpathy/rendergit) but for local directories.
"""

import os
import sys
import argparse
from pathlib import Path
from typing import List, Dict, Set
import mimetypes
import re
import fnmatch

# File extensions to skip (binary files, etc.)
SKIP_EXTENSIONS = {
    '.pyc', '.pyo', '.pyd', '.so', '.dll', '.exe', '.bin', '.obj', '.o',
    '.a', '.lib', '.dylib', '.class', '.jar', '.war', '.ear',
    '.zip', '.tar', '.gz', '.bz2', '.xz', '.7z', '.rar',
    '.jpg', '.jpeg', '.png', '.gif', '.bmp', '.ico', '.svg',
    '.mp3', '.mp4', '.avi', '.mov', '.wmv', '.flv', '.mkv',
    '.pdf', '.doc', '.docx', '.xls', '.xlsx', '.ppt', '.pptx',
    '.db', '.sqlite', '.sqlite3', '.mdb', '.accdb',
    '.git', '.gitignore', '.DS_Store', 'Thumbs.db'
}

# Directories to skip
SKIP_DIRS = {
    '__pycache__', '.git', '.svn', '.hg', '.bzr',
    'node_modules', '.venv', 'venv', 'env', '.env',
    'build', 'dist', 'target', 'bin', 'obj',
    '.idea', '.vscode', '.vs'
}

# Maximum file size to include (in bytes) - 1MB
MAX_FILE_SIZE = 1024 * 1024

def parse_gitignore_patterns(gitignore_file: Path) -> List[str]:
    """Parse gitignore patterns from a file."""
    patterns = []
    try:
        with open(gitignore_file, 'r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                # Skip empty lines and comments
                if line and not line.startswith('#'):
                    patterns.append(line)
    except Exception as e:
        print(f"Warning: Could not read gitignore file {gitignore_file}: {e}")
    return patterns

def matches_gitignore_pattern(path: Path, patterns: List[str], base_dir: Path) -> bool:
    """Check if a path matches any gitignore pattern."""
    if not patterns:
        return False
    
    # Convert path to relative path from base directory
    try:
        relative_path = path.relative_to(base_dir)
        path_str = str(relative_path).replace('\\', '/')  # Normalize separators
    except ValueError:
        # Path is not relative to base directory
        return False
    
    for pattern in patterns:
        # Handle different pattern types
        if pattern.startswith('/'):
            # Pattern starts with / - matches from root
            if fnmatch.fnmatch(path_str, pattern[1:]):
                return True
        elif pattern.endswith('/'):
            # Pattern ends with / - matches directories
            if path.is_dir() and fnmatch.fnmatch(path_str, pattern[:-1]):
                return True
        else:
            # Regular pattern - matches anywhere in path
            if fnmatch.fnmatch(path_str, pattern) or fnmatch.fnmatch(path_str, f"*/{pattern}"):
                return True
    
    return False

def is_binary_file(file_path: Path) -> bool:
    """Check if a file is binary by reading the first few bytes."""
    try:
        with open(file_path, 'rb') as f:
            chunk = f.read(1024)
            return b'\x00' in chunk
    except Exception:
        return True

def should_skip_file(file_path: Path, gitignore_patterns: List[str] = None, base_dir: Path = None) -> bool:
    """Determine if a file should be skipped."""
    # Skip if extension is in skip list
    if file_path.suffix.lower() in SKIP_EXTENSIONS:
        return True
    
    # Skip if file is too large
    try:
        if file_path.stat().st_size > MAX_FILE_SIZE:
            return True
    except Exception:
        return True
    
    # Skip if it's a binary file
    if is_binary_file(file_path):
        return True
    
    # Skip if it matches gitignore patterns
    if gitignore_patterns and base_dir:
        if matches_gitignore_pattern(file_path, gitignore_patterns, base_dir):
            return True
    
    return False

def should_skip_directory(dir_path: Path, gitignore_patterns: List[str] = None, base_dir: Path = None) -> bool:
    """Determine if a directory should be skipped."""
    # Check built-in skip list
    if dir_path.name in SKIP_DIRS:
        return True
    
    # Check gitignore patterns
    if gitignore_patterns and base_dir:
        if matches_gitignore_pattern(dir_path, gitignore_patterns, base_dir):
            return True
    
    return False

def get_file_content(file_path: Path) -> str:
    """Get the content of a file as a string."""
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            return f.read()
    except UnicodeDecodeError:
        try:
            with open(file_path, 'r', encoding='latin-1') as f:
                return f.read()
        except Exception:
            return f"[Error: Could not read file {file_path}]"
    except Exception as e:
        return f"[Error reading {file_path}: {str(e)}]"

def collect_files(directory: Path, gitignore_patterns: List[str] = None) -> List[Dict]:
    """Collect all files from the directory recursively."""
    files = []
    
    for root, dirs, files_in_dir in os.walk(directory):
        root_path = Path(root)
        
        # Remove directories that should be skipped
        dirs[:] = [d for d in dirs if not should_skip_directory(root_path / d, gitignore_patterns, directory)]
        
        for file_name in files_in_dir:
            file_path = root_path / file_name
            
            if should_skip_file(file_path, gitignore_patterns, directory):
                continue
            
            try:
                relative_path = file_path.relative_to(directory)
                file_size = file_path.stat().st_size
                content = get_file_content(file_path)
                
                files.append({
                    'path': relative_path,
                    'size': file_size,
                    'content': content,
                    'extension': file_path.suffix.lower()
                })
            except Exception as e:
                print(f"Warning: Could not process {file_path}: {e}")
    
    return sorted(files, key=lambda x: str(x['path']))

def generate_output_content(files: List[Dict], directory_name: str) -> str:
    """Generate the flattened content."""
    output_lines = []
    
    # Header
    output_lines.append(f"# Flattened Directory: {directory_name}")
    output_lines.append(f"# Generated on: {os.popen('date').read().strip()}")
    output_lines.append(f"# Total files: {len(files)}")
    output_lines.append("")
    
    # Directory tree overview
    output_lines.append("## Directory Structure")
    output_lines.append("")
    
    current_dir = ""
    for file_info in files:
        path_parts = str(file_info['path']).split('/')
        if len(path_parts) > 1:
            for i, part in enumerate(path_parts[:-1]):
                indent = "  " * i
                if current_dir != "/".join(path_parts[:i+1]):
                    output_lines.append(f"{indent}📁 {part}/")
                    current_dir = "/".join(path_parts[:i+1])
        
        indent = "  " * (len(path_parts) - 1)
        size_str = f"({file_info['size']} bytes)" if file_info['size'] < 1024 else f"({file_info['size'] // 1024} KB)"
        output_lines.append(f"{indent}📄 {path_parts[-1]} {size_str}")
    
    output_lines.append("")
    output_lines.append("---")
    output_lines.append("")
    
    # File contents
    for i, file_info in enumerate(files, 1):
        output_lines.append(f"## File {i}: {file_info['path']}")
        output_lines.append(f"Size: {file_info['size']} bytes")
        output_lines.append(f"Extension: {file_info['extension']}")
        output_lines.append("")
        output_lines.append("```")
        
        # Add language hint for syntax highlighting if possible
        lang_map = {
            '.py': 'python',
            '.js': 'javascript',
            '.ts': 'typescript',
            '.html': 'html',
            '.css': 'css',
            '.scss': 'scss',
            '.sass': 'sass',
            '.json': 'json',
            '.xml': 'xml',
            '.yaml': 'yaml',
            '.yml': 'yaml',
            '.toml': 'toml',
            '.ini': 'ini',
            '.cfg': 'ini',
            '.conf': 'ini',
            '.md': 'markdown',
            '.txt': 'text',
            '.sh': 'bash',
            '.bash': 'bash',
            '.zsh': 'bash',
            '.fish': 'bash',
            '.c': 'c',
            '.cpp': 'cpp',
            '.cc': 'cpp',
            '.cxx': 'cpp',
            '.h': 'c',
            '.hpp': 'cpp',
            '.java': 'java',
            '.rb': 'ruby',
            '.php': 'php',
            '.go': 'go',
            '.rs': 'rust',
            '.swift': 'swift',
            '.kt': 'kotlin',
            '.scala': 'scala',
            '.sql': 'sql',
            '.r': 'r',
            '.m': 'matlab',
            '.jl': 'julia',
            '.dart': 'dart',
            '.lua': 'lua',
            '.pl': 'perl',
            '.tcl': 'tcl',
            '.v': 'verilog',
            '.vhdl': 'vhdl',
            '.f': 'fortran',
            '.f90': 'fortran',
            '.f95': 'fortran',
            '.hs': 'haskell',
            '.ml': 'ocaml',
            '.fs': 'fsharp',
            '.clj': 'clojure',
            '.scm': 'scheme',
            '.lisp': 'lisp',
            '.el': 'elisp',
            '.vim': 'vim',
            '.tex': 'latex',
            '.bib': 'bibtex',
            '.dockerfile': 'dockerfile',
            '.makefile': 'makefile',
            '.cmake': 'cmake',
            '.gradle': 'gradle',
            '.pom': 'xml',
            '.sbt': 'scala',
            '.csproj': 'xml',
            '.vcxproj': 'xml',
            '.sln': 'text',
            '.gitignore': 'gitignore',
            '.gitattributes': 'gitattributes',
            '.editorconfig': 'ini',
            '.eslintrc': 'json',
            '.prettierrc': 'json',
            '.babelrc': 'json',
            '.webpack': 'javascript',
            '.rollup': 'javascript',
            '.vite': 'javascript',
            '.jest': 'javascript',
            '.mocha': 'javascript',
            '.karma': 'javascript',
            '.travis': 'yaml',
            '.github': 'yaml',
            '.gitlab': 'yaml',
            '.circleci': 'yaml',
            '.azure': 'yaml',
            '.aws': 'yaml',
            '.terraform': 'hcl',
            '.tf': 'hcl',
            '.tfvars': 'hcl',
            '.helm': 'yaml',
            '.kubernetes': 'yaml',
            '.k8s': 'yaml',
            '.docker': 'dockerfile',
            '.dockerignore': 'gitignore',
            '.env': 'bash',
            '.env.local': 'bash',
            '.env.development': 'bash',
            '.env.production': 'bash',
            '.env.test': 'bash',
            '.npmrc': 'ini',
            '.yarnrc': 'ini',
            '.bowerrc': 'json',
            '.composer': 'json',
            '.pip': 'ini',
            '.conda': 'yaml',
            '.requirements': 'text',
            '.setup': 'python',
            '.pyproject': 'toml',
            '.poetry': 'toml',
            '.cargo': 'toml',
            '.crates': 'toml',
            '.go.mod': 'go',
            '.go.sum': 'text',
            '.package.json': 'json',
            '.package-lock.json': 'json',
            '.yarn.lock': 'text',
            '.pnpm-lock.yaml': 'yaml',
            '.Cargo.lock': 'toml',
            '.Gemfile': 'ruby',
            '.Gemfile.lock': 'text',
            '.composer.json': 'json',
            '.composer.lock': 'json',
            '.pom.xml': 'xml',
            '.build.gradle': 'gradle',
            '.build.sbt': 'scala',
            '.project': 'xml',
            '.classpath': 'xml',
            '.settings': 'xml',
            '.launch': 'xml',
            '.tasks': 'json',
            '.vscode': 'json',
            '.idea': 'xml',
            '.xcodeproj': 'text',
            '.xcworkspace': 'text',
            '.pbxproj': 'text',
            '.storyboard': 'xml',
            '.xib': 'xml',
            '.plist': 'xml',
            '.entitlements': 'xml',
            '.provisionprofile': 'text',
            '.mobileprovision': 'text',
            '.p12': 'text',
            '.cer': 'text',
            '.pem': 'text',
            '.key': 'text',
            '.crt': 'text',
            '.csr': 'text',
            '.p7b': 'text',
            '.pfx': 'text',
            '.jks': 'text',
            '.keystore': 'text',
            '.truststore': 'text',
            '.cacerts': 'text',
            '.pem': 'text',
            '.key': 'text',
            '.crt': 'text',
            '.csr': 'text',
            '.p7b': 'text',
            '.pfx': 'text',
            '.jks': 'text',
            '.keystore': 'text',
            '.truststore': 'text',
            '.cacerts': 'text'
        }
        
        lang = lang_map.get(file_info['extension'], '')
        if lang:
            output_lines.append(f"```{lang}")
        else:
            output_lines.append("```")
        
        output_lines.append(file_info['content'])
        output_lines.append("```")
        output_lines.append("")
        output_lines.append("---")
        output_lines.append("")
    
    return "\n".join(output_lines)

def main():
    global MAX_FILE_SIZE, SKIP_EXTENSIONS, SKIP_DIRS
    
    parser = argparse.ArgumentParser(
        description="Flatten all files in a directory into a single file",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python flatten_directory.py /path/to/directory
  python flatten_directory.py /path/to/directory -o output.txt
  python flatten_directory.py /path/to/directory --output flattened_code.md
  python flatten_directory.py /path/to/directory --gitignore .flattenignore
  python flatten_directory.py /path/to/directory --verbose
        """
    )
    
    parser.add_argument(
        'directory',
        help='Directory to flatten'
    )
    
    parser.add_argument(
        '-o', '--output',
        help='Output file path (default: flattened_<directory_name>.txt)'
    )
    
    parser.add_argument(
        '--max-size',
        type=int,
        default=MAX_FILE_SIZE,
        help=f'Maximum file size to include in bytes (default: {MAX_FILE_SIZE})'
    )
    
    parser.add_argument(
        '--skip-extensions',
        nargs='+',
        help='Additional file extensions to skip'
    )
    
    parser.add_argument(
        '--skip-dirs',
        nargs='+',
        help='Additional directories to skip'
    )
    
    parser.add_argument(
        '--gitignore',
        help='Path to gitignore-style file with patterns to skip (e.g., .flattenignore)'
    )
    
    parser.add_argument(
        '--verbose', '-v',
        action='store_true',
        help='Show detailed information including list of files to be flattened'
    )
    
    args = parser.parse_args()
    
    # Update global settings based on arguments
    MAX_FILE_SIZE = args.max_size
    
    if args.skip_extensions:
        SKIP_EXTENSIONS.update({ext.lower() if not ext.startswith('.') else ext.lower() for ext in args.skip_extensions})
    
    if args.skip_dirs:
        SKIP_DIRS.update(args.skip_dirs)
    
    # Parse gitignore patterns
    gitignore_patterns = []
    if args.gitignore:
        gitignore_file = Path(args.gitignore)
        if gitignore_file.exists():
            gitignore_patterns = parse_gitignore_patterns(gitignore_file)
            print(f"Loaded {len(gitignore_patterns)} patterns from {gitignore_file}")
        else:
            print(f"Warning: Gitignore file {gitignore_file} not found")
    
    # Validate directory
    directory = Path(args.directory)
    if not directory.exists():
        print(f"Error: Directory '{directory}' does not exist.")
        sys.exit(1)
    
    if not directory.is_dir():
        print(f"Error: '{directory}' is not a directory.")
        sys.exit(1)
    
    print(f"Scanning directory: {directory}")
    
    # Collect files
    files = collect_files(directory, gitignore_patterns)
    
    if not files:
        print("No files found to flatten.")
        sys.exit(0)
    
    print(f"Found {len(files)} files to flatten.")
    
    # Show detailed file list if verbose mode is enabled
    if args.verbose:
        print("\n📁 Files to be flattened:")
        print("=" * 50)
        for i, file_info in enumerate(files, 1):
            size_str = f"({file_info['size']} bytes)" if file_info['size'] < 1024 else f"({file_info['size'] // 1024} KB)"
            print(f"{i:3d}. {file_info['path']} {size_str}")
        print("=" * 50)
        print()
    
    # Generate output filename
    if args.output:
        output_file = Path(args.output)
    else:
        output_file = Path(f"flattened_{directory.name}.txt")
    
    # Generate content
    content = generate_output_content(files, directory.name)
    
    # Write to file
    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(content)
        print(f"Successfully flattened directory to: {output_file}")
        print(f"Total files processed: {len(files)}")
        print(f"Output file size: {len(content)} characters")
    except Exception as e:
        print(f"Error writing output file: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()