nixos-migration_v9.py

#!/usr/bin/env python3

import os
import re
import shutil
import argparse
import subprocess
import time
from pathlib import Path

class PackageMigrationError(Exception):
    """Custom exception for migration errors"""
    pass

class GitError(Exception):
    """Exception for Git errors"""
    pass

def clean_empty_lines(content):
    """Clean consecutive empty lines, keeping only one"""
    # Split content into lines
    lines = content.splitlines()
    
    # Remove trailing whitespace from each line
    lines = [line.rstrip() for line in lines]
    
    # Remove consecutive empty lines
    cleaned_lines = []
    prev_empty = False
    for line in lines:
        if not line.strip():  # Empty line
            if not prev_empty:
                cleaned_lines.append(line)
            prev_empty = True
        else:
            cleaned_lines.append(line)
            prev_empty = False
            
    # Ensure file ends with a single newline
    return '\n'.join(cleaned_lines) + '\n'

def setup_git_branch(nixpkgs_path):
    """Setup a working branch for migrations"""
    try:
        current_branch = subprocess.run(
            ['git', 'rev-parse', '--abbrev-ref', 'HEAD'],
            cwd=nixpkgs_path,
            capture_output=True,
            text=True,
            check=True
        ).stdout.strip()
        
        # Create and switch to a new branch
        branch_name = f"by-name-migration-{int(time.time())}"
        subprocess.run(['git', 'checkout', '-b', branch_name], cwd=nixpkgs_path, check=True)
        print(f"Created and switched to branch: {branch_name}")
        return current_branch
    except subprocess.CalledProcessError as e:
        raise GitError(f"Git error: {str(e)}")

def check_custom_arguments(package_content, original_line):
    """Checks if the package uses custom arguments in all-packages.nix"""
    return '{' in original_line and '}' in original_line and not original_line.strip().endswith('{ };')

def run_nixpkgs_vet(nixpkgs_path):
    """
    Run nixpkgs-vet checks on the repository
    Returns (bool, str) - Success status and error message if any
    """
    try:
        result = subprocess.run(
            ['./ci/nixpkgs-vet.sh', 'master', 'https://github.com/NixOS/nixpkgs.git'],
            cwd=nixpkgs_path,
            capture_output=True,
            text=True,
            check=True
        )
        return True, ""
    except subprocess.CalledProcessError as e:
        return False, e.stdout + e.stderr

def verify_package_files(package_dir):
    """
    Verify that all files referenced in package.nix exist
    Returns (bool, list of missing files)
    """
    package_nix = os.path.join(package_dir, 'package.nix')
    if not os.path.exists(package_nix):
        return False, ["package.nix"]
        
    missing_files = []
    with open(package_nix, 'r') as f:
        content = f.read()
        
    # Find all file references in quotes
    file_refs = re.findall(r'["\']\./[^"\']+["\']', content)
    for ref in file_refs:
        # Remove quotes and ./ from the beginning
        file_path = ref.strip('\'"').replace('./', '')
        full_path = os.path.join(package_dir, file_path)
        if not os.path.exists(full_path):
            missing_files.append(file_path)
            
    return len(missing_files) == 0, missing_files

def git_commit_changes(nixpkgs_path, package_name):
    """Creates a Git commit for package migration"""
    try:
        os.chdir(nixpkgs_path)
        
        # Stash any unrelated changes
        subprocess.run(['git', 'stash', '--keep-index'], check=False)
        
        # Add only specific directories and files
        subprocess.run(['git', 'add', 'pkgs/by-name'], check=True)
        subprocess.run(['git', 'add', 'pkgs/top-level/all-packages.nix'], check=True)
        subprocess.run(['git', 'add', 'pkgs/applications'], check=True)
        subprocess.run(['git', 'add', 'pkgs/tools'], check=True)
        subprocess.run(['git', 'add', 'pkgs/editors'], check=True)  # Add editors directory for xxe-pe
        
        # Create commit
        commit_message = f"{package_name}: move to by-name"
        subprocess.run(['git', 'commit', '-m', commit_message], check=True)
        print(f"Created commit: {commit_message}")
        
    except subprocess.CalledProcessError as e:
        raise GitError(f"Git error: {str(e)}")

def find_packages_to_migrate(nixpkgs_path, source_dir):
    """Finds all packages eligible for migration from specified source directory"""
    all_packages_path = os.path.join(nixpkgs_path, 'pkgs/top-level/all-packages.nix')
    packages = []
    
    with open(all_packages_path, 'r') as f:
        content = f.readlines()
    
    # Create pattern based on source directory
    source_pattern = source_dir.replace('/', '\/')
    patterns = [
        rf'([\w-]+)\s*=\s*callPackage\s+\.\./{source_pattern}\/([^\s{{}}]+)\s*{{\s*}};',  # Basic format with hyphens
        rf'([\w-]+)\s*=\s*callPackage\s+\.\./{source_pattern}\/([^\s{{}}]+\/[^\s{{}}]+)\s*{{\s*}};',  # With subdirectory
    ]
    
    for line in content:
        line = line.strip()
        if not line or line.startswith('#'):
            continue
            
        for pattern in patterns:
            match = re.search(pattern, line)
            if match:
                package_name = match.group(1)  # Full package name including hyphens
                path_part = match.group(2)
                path = f"../{source_dir}/{path_part}"
                packages.append((package_name, path, line))
                break
    
    return packages

def migrate_package(nixpkgs_path, package_name, source_path, original_line, dry_run=False, auto_commit=False, skip_vet=False):
    """Migrates a package to the by-name structure"""
    dest_dir = None
    try:
        os.chdir(nixpkgs_path)
        all_packages_path = 'pkgs/top-level/all-packages.nix'
        
        with open(all_packages_path, 'r') as f:
            content = f.read()
            
        has_custom_args = check_custom_arguments(content, original_line)
        if has_custom_args:
            print(f"WARNING: {package_name} has custom arguments in all-packages.nix")
            print(f"Definition line: {original_line}")
            print("Definition in all-packages.nix must be kept")
            
        # Get first two letters considering possible prefixes
        if '-' in package_name:
            prefix = package_name.split('-')[0]
            first_two_letters = prefix[:2].lower()
        else:
            first_two_letters = package_name[:2].lower()
        
        dest_dir = os.path.join('pkgs', 'by-name', first_two_letters, package_name)
        
        if dry_run:
            print(f"[DRY RUN] Would migrate: {package_name}")
            if not has_custom_args:
                print(f"Would remove line: {original_line}")
            return True
        
        os.makedirs(os.path.join('pkgs', 'by-name', first_two_letters), exist_ok=True)
        
        # Fix path handling to get directory path
        normalized_path = source_path.replace('../', '')
        if normalized_path.endswith('default.nix'):
            source_dir = os.path.dirname(os.path.join('pkgs', normalized_path))
        else:
            source_dir = os.path.join('pkgs', normalized_path)
        
        if not os.path.exists(source_dir) or not os.path.isdir(source_dir):
            # Try alternative path formats
            alt_paths = [
                os.path.dirname(os.path.join('pkgs', normalized_path)),
                os.path.join('pkgs', source_dir.split('/')[-2], package_name),  # Use category from path
                os.path.dirname(os.path.join('pkgs', source_path))
            ]
            
            for alt_path in alt_paths:
                if os.path.exists(alt_path) and os.path.isdir(alt_path):
                    source_dir = alt_path
                    break
            else:
                raise PackageMigrationError(f"Source directory not found. Tried:\n" + 
                                          "\n".join([f"- {p}" for p in [source_dir] + alt_paths]))
        
        print(f"Source directory: {source_dir}")
        
        # First, copy all files
        os.makedirs(dest_dir, exist_ok=True)
        for item in os.listdir(source_dir):
            src_item = os.path.join(source_dir, item)
            dst_item = os.path.join(dest_dir, item)
            if os.path.isfile(src_item):
                shutil.copy2(src_item, dst_item)
            elif os.path.isdir(src_item):
                shutil.copytree(src_item, dst_item)
        
        # Then rename default.nix to package.nix
        if os.path.exists(os.path.join(dest_dir, 'default.nix')):
            os.rename(
                os.path.join(dest_dir, 'default.nix'),
                os.path.join(dest_dir, 'package.nix')
            )
        
        # Verify all required files exist
        files_ok, missing_files = verify_package_files(dest_dir)
        if not files_ok:
            # Try to copy missing files
            src_files_found = False
            potential_paths = [
                source_dir,
                os.path.dirname(source_dir),
                os.path.join(source_dir, '..')
            ]
            
            for missing_file in missing_files:
                for path in potential_paths:
                    src_file = os.path.join(path, missing_file)
                    if os.path.exists(src_file):
                        dst_file = os.path.join(dest_dir, missing_file)
                        os.makedirs(os.path.dirname(dst_file), exist_ok=True)
                        shutil.copy2(src_file, dst_file)
                        src_files_found = True
                        print(f"Copied missing file: {missing_file}")
                        break
                        
            if not src_files_found:
                raise PackageMigrationError(
                    f"Missing required files for {package_name}: {', '.join(missing_files)}"
                )
        
        if not has_custom_args:
            with open(all_packages_path, 'r') as f:
                content = f.read()
            
            # Remove the package line and clean consecutive empty lines
            new_content = content.replace(original_line + '\n', '')
            new_content = clean_empty_lines(new_content)
            
            with open(all_packages_path, 'w') as f:
                f.write(new_content)
            print(f"Removed line from all-packages.nix: {original_line}")
            print("Cleaned empty lines in all-packages.nix")
        
        # Remove source directory
        shutil.rmtree(source_dir)
        print(f"Removed source directory: {source_dir}")
        
        if auto_commit:
            # Run nixpkgs-vet before committing
            if not skip_vet:
                print("Running nixpkgs-vet check...")
                vet_success, vet_output = run_nixpkgs_vet(nixpkgs_path)
                if not vet_success:
                    print("nixpkgs-vet check failed:")
                    print(vet_output)
                    print(f"Skipping {package_name} due to nixpkgs-vet check failure")
                    return False
                print("nixpkgs-vet check passed")
            
            git_commit_changes(nixpkgs_path, package_name)
        
        print(f"Successfully migrated {package_name}")
        return True
        
    except (PackageMigrationError, GitError) as e:
        print(f"Error migrating {package_name}: {str(e)}")
        return False

def main():
    parser = argparse.ArgumentParser(description='Migrate NixOS packages to by-name structure')
    parser.add_argument('nixpkgs_path', help='Path to nixpkgs repository')
    parser.add_argument('--source-dir', default='applications/misc',
                      help='Source directory path relative to pkgs (e.g., applications/misc, tools/audio)')
    parser.add_argument('--dry-run', action='store_true', help='Simulate migration without making changes')
    parser.add_argument('--package', help='Migrate a specific package')
    parser.add_argument('--auto-commit', action='store_true', help='Automatically create Git commits')
    parser.add_argument('--skip-vet', action='store_true', help='Skip nixpkgs-vet check')
    args = parser.parse_args()
    
    nixpkgs_path = os.path.abspath(args.nixpkgs_path)
    
    if not os.path.exists(nixpkgs_path):
        print(f"Error: Directory {nixpkgs_path} does not exist")
        return
        
    if not os.path.exists(os.path.join(nixpkgs_path, '.git')):
        print(f"Error: {nixpkgs_path} is not a Git repository")
        return
    
    original_dir = os.getcwd()
    original_branch = None
    
    try:
        if args.auto_commit:
            original_branch = setup_git_branch(nixpkgs_path)
        
        if args.package:
            packages = [(p, path, line) for p, path, line in find_packages_to_migrate(nixpkgs_path, args.source_dir) 
                       if p == args.package]
            if not packages:
                print(f"Package {args.package} not found or not eligible for migration")
                return
        else:
            packages = find_packages_to_migrate(nixpkgs_path, args.source_dir)
        
        print(f"{'[DRY RUN] ' if args.dry_run else ''}Found {len(packages)} package(s) to migrate from {args.source_dir}")
        
        for package_name, source_path, original_line in packages:
            print(f"\nMigrating {package_name}...")
            success = migrate_package(nixpkgs_path, package_name, source_path, original_line, args.dry_run, args.auto_commit, args.skip_vet)
            if not success:
                print(f"Migration failed for {package_name}, but continuing with next package...")
            
    finally:
        os.chdir(original_dir)
        if original_branch and args.auto_commit:
            try:
                subprocess.run(['git', 'checkout', original_branch], cwd=nixpkgs_path, check=True)
                print(f"Switched back to branch {original_branch}")
            except:
                print(f"Warning: Could not switch back to branch {original_branch}")

if __name__ == "__main__":
    main()