scripts/utilities/standardize_imports.py

#!/usr/bin/env python3
"""
Import Standardization Script

This script standardizes Python imports across the codebase according to PEP 8:
1. Standard library imports first
2. Third-party imports second  
3. Local application imports third
4. Remove unused imports
5. Sort imports within each group alphabetically
"""

import ast
from collections import defaultdict
import os
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple

from dataclasses import dataclass

# Standard library modules (Python 3.8+)
STDLIB_MODULES = {
    'abc', 'argparse', 'ast', 'asyncio', 'base64', 'collections', 'copy', 
    'datetime', 'decimal', 'enum', 'functools', 'hashlib', 'io', 'itertools',
    'json', 'logging', 'math', 'os', 'pathlib', 'pickle', 're', 'shutil',
    'ssl', 'string', 'subprocess', 'sys', 'tempfile', 'threading', 'time',
    'typing', 'uuid', 'warnings', 'weakref', 'certifi'
}

# Third-party modules commonly used in this project
THIRD_PARTY_MODULES = {
    'fastapi', 'pydantic', 'httpx', 'pytest', 'motor', 'pymongo', 'nltk',
    'spacy', 'rake_nltk', 'duckduckgo_search', 'google', 'genai', 'dotenv'
}

@dataclass
class ImportInfo:
    """Information about an import statement"""
    module: str
    names: List[str]
    alias: Optional[str]
    is_from_import: bool
    line_number: int
    original_line: str

class ImportStandardizer:
    """Standardizes imports in Python files"""
    
    def __init__(self):
        self.used_names: Set[str] = set()
        self.imports: List[ImportInfo] = []
        
    def analyze_file(self, file_path: str) -> Tuple[List[ImportInfo], Set[str]]:
        """Analyze a Python file to extract imports and used names"""
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
            
        try:
            tree = ast.parse(content)
        except SyntaxError as e:
            print(f"Syntax error in {file_path}: {e}")
            return [], set()
            
        imports = []
        used_names = set()
        
        # Extract imports
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for alias in node.names:
                    imports.append(ImportInfo(
                        module=alias.name,
                        names=[],
                        alias=alias.asname,
                        is_from_import=False,
                        line_number=node.lineno,
                        original_line=self._get_line_from_content(content, node.lineno)
                    ))
                    
            elif isinstance(node, ast.ImportFrom):
                if node.module:  # Skip relative imports without module
                    names = [alias.name for alias in node.names]
                    imports.append(ImportInfo(
                        module=node.module,
                        names=names,
                        alias=None,
                        is_from_import=True,
                        line_number=node.lineno,
                        original_line=self._get_line_from_content(content, node.lineno)
                    ))
                    
            # Extract used names (simplified - looks for Name nodes)
            elif isinstance(node, ast.Name):
                used_names.add(node.id)
                
        return imports, used_names
        
    def _get_line_from_content(self, content: str, line_number: int) -> str:
        """Get the actual line content from file"""
        lines = content.split('\n')
        if 1 <= line_number <= len(lines):
            return lines[line_number - 1].strip()
        return ""
        
    def categorize_import(self, import_info: ImportInfo) -> str:
        """Categorize import as stdlib, third-party, or local"""
        module_root = import_info.module.split('.')[0]
        
        if module_root in STDLIB_MODULES:
            return 'stdlib'
        elif module_root in THIRD_PARTY_MODULES:
            return 'third_party'
        elif module_root in ['analytics', 'tests', 'scripts']:
            return 'local'
        else:
            # Default to third-party for unknown modules
            return 'third_party'
            
    def is_import_used(self, import_info: ImportInfo, used_names: Set[str]) -> bool:
        """Check if an import is actually used in the code"""
        if import_info.is_from_import:
            # For 'from module import name', check if any imported names are used
            return any(name in used_names for name in import_info.names)
        else:
            # For 'import module', check if module name or alias is used
            check_name = import_info.alias if import_info.alias else import_info.module.split('.')[0]
            return check_name in used_names
            
    def generate_standardized_imports(self, imports: List[ImportInfo], used_names: Set[str]) -> str:
        """Generate standardized import section"""
        # Filter out unused imports
        used_imports = [imp for imp in imports if self.is_import_used(imp, used_names)]
        
        # Categorize imports
        categorized = defaultdict(list)
        for imp in used_imports:
            category = self.categorize_import(imp)
            categorized[category].append(imp)
            
        # Sort imports within each category
        for category in categorized:
            categorized[category].sort(key=lambda x: (x.module, x.is_from_import))
            
        # Generate import strings
        import_lines = []
        
        # Standard library imports
        if categorized['stdlib']:
            for imp in categorized['stdlib']:
                import_lines.append(self._format_import(imp))
            import_lines.append('')  # Empty line after stdlib
            
        # Third-party imports
        if categorized['third_party']:
            for imp in categorized['third_party']:
                import_lines.append(self._format_import(imp))
            import_lines.append('')  # Empty line after third-party
            
        # Local imports
        if categorized['local']:
            for imp in categorized['local']:
                import_lines.append(self._format_import(imp))
            import_lines.append('')  # Empty line after local
                
        return '\n'.join(import_lines).rstrip() + '\n'
        
    def _format_import(self, import_info: ImportInfo) -> str:
        """Format an import statement"""
        if import_info.is_from_import:
            names_str = ', '.join(sorted(import_info.names))
            return f"from {import_info.module} import {names_str}"
        else:
            if import_info.alias:
                return f"import {import_info.module} as {import_info.alias}"
            else:
                return f"import {import_info.module}"
                
    def standardize_file(self, file_path: str) -> bool:
        """Standardize imports in a single file"""
        print(f"Standardizing imports in: {file_path}")
        
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
            
        imports, used_names = self.analyze_file(file_path)
        
        if not imports:
            print(f"  No imports found in {file_path}")
            return False
            
        # Generate standardized imports
        standardized_imports = self.generate_standardized_imports(imports, used_names)
        
        # More careful replacement - find import blocks and preserve structure
        lines = content.split('\n')
        
        # Find all import line numbers
        import_lines = set(imp.line_number - 1 for imp in imports)  # Convert to 0-based
        
        # Find contiguous import blocks
        import_blocks = []
        current_block = []
        
        for i, line in enumerate(lines):
            if i in import_lines:
                current_block.append(i)
            elif current_block:
                import_blocks.append(current_block)
                current_block = []
        
        if current_block:
            import_blocks.append(current_block)
            
        # Only replace if we have a single contiguous import block
        if len(import_blocks) == 1:
            block = import_blocks[0]
            first_line = block[0]
            last_line = block[-1]
            
            # Find the end of the import section (including empty lines)
            end_line = last_line
            for i in range(last_line + 1, len(lines)):
                line = lines[i].strip()
                if not line:  # Empty line
                    end_line = i
                elif line.startswith('#'):  # Comment
                    continue
                else:
                    break
                    
            # Replace the import block
            new_lines = (
                lines[:first_line] +
                standardized_imports.rstrip().split('\n') +
                lines[end_line + 1:]
            )
            
            new_content = '\n'.join(new_lines)
            
            # Validate the new content can be parsed
            try:
                ast.parse(new_content)
            except SyntaxError as e:
                print(f"  ⚠️  Syntax error after standardization, skipping: {e}")
                return False
            
            # Write back to file
            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(new_content)
                
            print(f"  ✓ Standardized {len(imports)} imports")
            return True
        else:
            print(f"  ⚠️  Multiple import blocks found, skipping for safety")
            return False
        
    def standardize_directory(self, directory: str, exclude_patterns: List[str] = None) -> int:
        """Standardize imports in all Python files in a directory"""
        if exclude_patterns is None:
            exclude_patterns = ['__pycache__', '.git', 'venv', 'env', '.pytest_cache']
            
        files_processed = 0
        
        for root, dirs, files in os.walk(directory):
            # Skip excluded directories
            dirs[:] = [d for d in dirs if not any(pattern in d for pattern in exclude_patterns)]
            
            for file in files:
                if file.endswith('.py'):
                    file_path = os.path.join(root, file)
                    try:
                        if self.standardize_file(file_path):
                            files_processed += 1
                    except Exception as e:
                        print(f"Error processing {file_path}: {e}")
                        
        return files_processed

def main():
    """Main function to standardize imports across the codebase"""
    print("🔧 Starting import standardization...")
    
    standardizer = ImportStandardizer()
    
    # Get the project root directory
    project_root = Path(__file__).parent.parent.parent
    
    # Directories to process (excluding backups)
    directories_to_process = [
        str(project_root / 'analytics'),
        str(project_root / 'tests'),
        str(project_root / 'scripts'),
        str(project_root / 'archive')
    ]
    
    # Process root level Python files individually
    root_files = [f for f in project_root.glob('*.py') if f.is_file()]
    
    total_files = 0
    
    # Process root level files
    print(f"\n📁 Processing root level Python files")
    for file_path in root_files:
        try:
            if standardizer.standardize_file(str(file_path)):
                total_files += 1
        except Exception as e:
            print(f"Error processing {file_path}: {e}")
    
    # Process directories (excluding backups)
    exclude_patterns = ['__pycache__', '.git', 'venv', 'env', '.pytest_cache', 'backups']
    
    for directory in directories_to_process:
        if os.path.exists(directory):
            print(f"\n📁 Processing directory: {directory}")
            files_processed = standardizer.standardize_directory(directory, exclude_patterns)
            total_files += files_processed
            print(f"  Processed {files_processed} files")
        else:
            print(f"⚠️  Directory not found: {directory}")
            
    print(f"\n✅ Import standardization complete!")
    print(f"📊 Total files processed: {total_files}")

if __name__ == "__main__":
    main()