result_cleaner.py

#!/usr/bin/env python3
"""
Result Cleaner - Clean all directories named 'result' or 'results' under edge-applications
"""

import os
import shutil
import argparse
from pathlib import Path


def find_result_directories(root_dir):
    """
    Recursively find all directories named 'result' or 'results'
    
    Args:
        root_dir: Root directory path
        
    Returns:
        list: List of paths to all result/results directories
    """
    result_dirs = []
    root_path = Path(root_dir)
    
    # Recursively traverse all subdirectories
    for dirpath, dirnames, _ in os.walk(root_path):
        for dirname in dirnames:
            if dirname in ['result', 'results']:
                full_path = os.path.join(dirpath, dirname)
                result_dirs.append(full_path)
    
    return result_dirs


def get_directory_size(path):
    """
    Calculate total size of a directory
    
    Args:
        path: Directory path
        
    Returns:
        int: Directory size in bytes
    """
    total_size = 0
    for dirpath, _, filenames in os.walk(path):
        for filename in filenames:
            filepath = os.path.join(dirpath, filename)
            try:
                # Use stat to avoid following symlinks and handle race conditions
                stat_info = os.lstat(filepath)
                # Only count regular files, not symlinks or special files
                if os.path.isfile(filepath) and not os.path.islink(filepath):
                    total_size += stat_info.st_size
            except (OSError, FileNotFoundError):
                # File might have been deleted or we don't have permission
                continue
    return total_size


def format_size(size_bytes):
    """
    Convert bytes to human-readable format
    
    Args:
        size_bytes: Size in bytes
        
    Returns:
        str: Formatted size string
    """
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if size_bytes < 1024.0:
            return f"{size_bytes:.2f} {unit}"
        size_bytes /= 1024.0
    return f"{size_bytes:.2f} PB"


def list_result_directories(root_dir):
    """
    List all found result/results directories
    
    Args:
        root_dir: Root directory path
    """
    result_dirs = find_result_directories(root_dir)
    
    if not result_dirs:
        print("No result or results directories found")
        return []
    
    print(f"\nFound {len(result_dirs)} result/results directories:\n")
    print(f"{'No.':<6} {'Path':<80} {'Size':<15}")
    print("-" * 101)
    
    total_size = 0
    for idx, dir_path in enumerate(result_dirs, 1):
        size = get_directory_size(dir_path)
        total_size += size
        # Get relative path for display
        rel_path = os.path.relpath(dir_path, root_dir)
        print(f"{idx:<6} {rel_path:<80} {format_size(size):<15}")
    
    print("-" * 101)
    print(f"{'Total:':<86} {format_size(total_size):<15}\n")
    
    return result_dirs


def clean_result_directories(root_dir, dry_run=False, confirm=True):
    """
    Clean all result/results directories
    
    Args:
        root_dir: Root directory path
        dry_run: If True, only show what would be deleted without actually deleting
        confirm: If True, require user confirmation before deletion
    """
    result_dirs = list_result_directories(root_dir)
    
    if not result_dirs:
        return
    
    if dry_run:
        print("[DRY RUN] The above directories would be deleted (simulation mode, not actually deleted)")
        return
    
    if confirm:
        response = input("Confirm deletion of all above directories? (yes/no): ").strip().lower()
        if response != 'yes':
            print("Operation cancelled")
            return
    
    print("Starting cleanup...")
    success_count = 0
    error_count = 0
    
    for dir_path in result_dirs:
        try:
            shutil.rmtree(dir_path)
            rel_path = os.path.relpath(dir_path, root_dir)
            print(f"✓ Deleted: {rel_path}")
            success_count += 1
        except Exception as e:
            rel_path = os.path.relpath(dir_path, root_dir)
            print(f"✗ Failed to delete: {rel_path} - {e}")
            error_count += 1
    
    print(f"\nCleanup complete!")
    print(f"Successfully deleted: {success_count} directories")
    if error_count > 0:
        print(f"Failed to delete: {error_count} directories")


def main():
    parser = argparse.ArgumentParser(
        description='Clean all directories named result or results under edge-applications',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s                     # Clean all result/results directories (default)
  %(prog)s --list              # Only list all result/results directories
  %(prog)s --dry-run           # Simulate run, show what would be deleted
        """
    )
    
    parser.add_argument(
        '--list',
        action='store_true',
        help='Only list all found result/results directories without deletion'
    )
    
    parser.add_argument(
        '--dry-run',
        action='store_true',
        help='Simulate run, show what would be deleted without actually deleting'
    )
    
    parser.add_argument(
        '--root-dir',
        type=str,
        default=None,
        help='Specify root directory (defaults to script directory)'
    )
    
    args = parser.parse_args()
    
    # Determine root directory
    if args.root_dir:
        root_dir = args.root_dir
    else:
        # Default to script directory
        root_dir = os.path.dirname(os.path.abspath(__file__))
    
    if not os.path.isdir(root_dir):
        print(f"Error: Directory does not exist: {root_dir}")
        return 1
    
    print(f"Searching directory: {root_dir}\n")
    
    # Execute operation based on arguments
    if args.list:
        list_result_directories(root_dir)
    elif args.dry_run:
        clean_result_directories(root_dir, dry_run=True, confirm=False)
    else:
        # Default behavior: clean without confirmation
        clean_result_directories(root_dir, dry_run=False, confirm=False)


if __name__ == '__main__':
    main()