#!/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()