update_to_memory_handler.py

#!/usr/bin/env python3
"""
Script to update all file upload handlers to use MemoryFileHandler
This will prevent 403 errors on restricted environments like Hugging Face Spaces
"""

import os
import re
from pathlib import Path

def update_file(filepath):
    """Update a single file to use MemoryFileHandler."""
    
    with open(filepath, 'r', encoding='utf-8') as f:
        content = f.read()
    
    original_content = content
    
    # Update import statements
    content = re.sub(
        r'from web_app\.utils import FileUploadHandler',
        'from web_app.utils import MemoryFileHandler',
        content
    )
    
    # Update FileUploadHandler calls to MemoryFileHandler
    content = re.sub(r'FileUploadHandler\.', 'MemoryFileHandler.', content)
    
    # Update save_to_temp pattern
    # Old pattern: temp_path = FileUploadHandler.save_to_temp(uploaded_file, prefix="...")
    # New pattern: content = MemoryFileHandler.process_uploaded_file(uploaded_file)
    
    # Replace save_to_temp and read_from_temp patterns
    content = re.sub(
        r'temp_path = MemoryFileHandler\.save_to_temp\(([^,]+), prefix="[^"]+"\)\s*'
        r'if temp_path:\s*'
        r'.*?= MemoryFileHandler\.read_from_temp\(temp_path\)',
        r'content = MemoryFileHandler.process_uploaded_file(\1)',
        content,
        flags=re.DOTALL
    )
    
    # Replace validate_file_size (MemoryFileHandler doesn't need this as it checks inline)
    content = re.sub(
        r'if not MemoryFileHandler\.validate_file_size\([^)]+\):\s*return',
        'if uploaded_file.size > 300 * 1024 * 1024:\n'
        '                        st.error(f"File too large ({uploaded_file.size / 1024 / 1024:.1f} MB). Maximum allowed: 300MB")\n'
        '                        return',
        content
    )
    
    # Remove cleanup_temp_file calls
    content = re.sub(
        r'MemoryFileHandler\.cleanup_temp_file\([^)]+\)\s*',
        '',
        content
    )
    
    # Remove cleanup_old_temp_files calls
    content = re.sub(
        r'MemoryFileHandler\.cleanup_old_temp_files\([^)]+\)\s*',
        '',
        content
    )
    
    if content != original_content:
        with open(filepath, 'w', encoding='utf-8') as f:
            f.write(content)
        print(f"✅ Updated: {filepath}")
        return True
    else:
        print(f"⏭️  No changes needed: {filepath}")
        return False

def main():
    """Main function to update all relevant files."""
    
    # Files to update
    files_to_update = [
        'web_app/handlers/analysis_handlers.py',
        'web_app/handlers/pos_handlers.py',
        'web_app/handlers/frequency_handlers.py',
        'web_app/reference_manager.py',
        'web_app/config_manager.py',
        'web_app/debug_utils.py'
    ]
    
    updated_count = 0
    
    print("🔄 Updating file handlers to use MemoryFileHandler...")
    print("=" * 60)
    
    for file_path in files_to_update:
        if os.path.exists(file_path):
            if update_file(file_path):
                updated_count += 1
        else:
            print(f"❌ File not found: {file_path}")
    
    print("=" * 60)
    print(f"✅ Updated {updated_count} files")
    
    # Create a backup of the old FileUploadHandler
    old_handler_path = 'web_app/utils/file_upload_handler.py'
    backup_path = 'web_app/utils/file_upload_handler.py.backup'
    
    if os.path.exists(old_handler_path) and not os.path.exists(backup_path):
        import shutil
        shutil.copy2(old_handler_path, backup_path)
        print(f"📋 Created backup: {backup_path}")

if __name__ == "__main__":
    main()