From 0d8cdd4193b8dc6b96c96495bcc2b8357529d3b6 Mon Sep 17 00:00:00 2001 From: psywarrior1998 Date: Sat, 9 Aug 2025 16:19:57 +0530 Subject: [PATCH] feat: Implement concurrent package upgrades --- upgrade_tool/main.py | 95 +++++++++++++++++++++++++++----------------- 1 file changed, 58 insertions(+), 37 deletions(-) diff --git a/upgrade_tool/main.py b/upgrade_tool/main.py index 57e20e0..dbba525 100644 --- a/upgrade_tool/main.py +++ b/upgrade_tool/main.py @@ -1,11 +1,14 @@ import subprocess import sys -from typing import List, Optional +from typing import List, Optional, Tuple import typer from rich.console import Console from rich.progress import Progress, SpinnerColumn, BarColumn, TextColumn +# Import the concurrent futures module for threading +from concurrent.futures import ThreadPoolExecutor, as_completed + # Import the refactored utility functions from .utils import get_outdated_packages, generate_packages_table @@ -19,6 +22,29 @@ add_completion=False, ) +def upgrade_package(pkg: dict) -> Tuple[str, str, bool]: + """ + Worker function to upgrade a single package in a separate thread. + + Args: + pkg: A dictionary containing package information ('name', 'latest_version'). + + Returns: + A tuple containing (package_name, latest_version, success_boolean). + """ + pkg_name = pkg['name'] + latest_version = pkg['latest_version'] + try: + # Execute the pip upgrade command, suppressing output + subprocess.check_call( + [sys.executable, "-m", "pip", "install", "--upgrade", pkg_name], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + return pkg_name, latest_version, True + except subprocess.CalledProcessError: + return pkg_name, latest_version, False + @app.command() def upgrade( packages_to_upgrade: Optional[List[str]] = typer.Argument( @@ -33,28 +59,27 @@ def upgrade( dry_run: bool = typer.Option( False, "--dry-run", help="Simulate the upgrade without making any changes." ), + workers: int = typer.Option( + 10, "--workers", "-w", help="Number of concurrent workers for parallel upgrades." + ) ): """ - Checks for and upgrades outdated Python packages. + Checks for and concurrently upgrades outdated Python packages. """ - # Use the utility function to get outdated packages + # --- Filtering Logic (Unchanged) --- outdated_packages = get_outdated_packages() if not outdated_packages: console.print("[bold green]✨ All packages are up to date! ✨[/bold green]") raise typer.Exit() - # --- Filtering Logic --- if packages_to_upgrade: - # User specified which packages to upgrade name_to_pkg = {pkg['name'].lower(): pkg for pkg in outdated_packages} target_packages = [name_to_pkg[name.lower()] for name in packages_to_upgrade if name.lower() in name_to_pkg] else: - # Default to all outdated packages target_packages = outdated_packages if exclude: - # Exclude packages specified by the user (case-insensitive) exclude_set = {name.lower() for name in exclude} target_packages = [pkg for pkg in target_packages if pkg['name'].lower() not in exclude_set] @@ -62,17 +87,15 @@ def upgrade( console.print("[bold yellow]No packages match the specified criteria for upgrade.[/bold yellow]") raise typer.Exit() - # --- Display and Confirmation --- - # Use the utility function to generate the table + # --- Display and Confirmation (Unchanged) --- table = generate_packages_table(target_packages, title="Outdated Python Packages") console.print(table) if dry_run: - console.print("\n[bold yellow]--dry-run enabled. No packages will be upgraded.[/bold yellow]") + console.print(f"\n[bold yellow]--dry-run enabled. Would upgrade {len(target_packages)} packages with {workers} workers.[/bold yellow]") raise typer.Exit() if not yes: - # Use Typer's confirmation prompt try: confirmed = typer.confirm("\nProceed with the upgrade?") if not confirmed: @@ -82,50 +105,48 @@ def upgrade( console.print("\nUpgrade cancelled by user.") raise typer.Exit() - # --- Execution Logic --- - console.print("\n[bold blue]Starting upgrade process...[/bold blue]") + # --- Concurrent Execution Logic (The New Engine) --- + console.print(f"\n[bold blue]Starting parallel upgrade with {workers} workers...[/bold blue]") - # Define a rich Progress bar progress = Progress( SpinnerColumn(), TextColumn("[progress.description]{task.description}"), BarColumn(), TextColumn("[progress.percentage]{task.percentage:>3.0f}%"), - console=console, + console=console ) + success_count = 0 + fail_count = 0 + with progress: upgrade_task = progress.add_task("[green]Upgrading...", total=len(target_packages)) - success_count = 0 - fail_count = 0 - for pkg in target_packages: - pkg_name = pkg['name'] - progress.update(upgrade_task, description=f"Upgrading [bold cyan]{pkg_name}[/bold cyan]...") + # Create a thread pool with the specified number of workers + with ThreadPoolExecutor(max_workers=workers) as executor: + # Submit an upgrade task for each package + future_to_pkg = {executor.submit(upgrade_package, pkg): pkg for pkg in target_packages} - try: - # Execute the pip upgrade command - subprocess.check_call( - [sys.executable, "-m", "pip", "install", "--upgrade", pkg_name], - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - ) - progress.console.print(f" ✅ [green]Successfully upgraded {pkg_name} to {pkg['latest_version']}[/green]") - success_count += 1 - except subprocess.CalledProcessError: - progress.console.print(f" ❌ [red]Failed to upgrade {pkg_name}[/red]") - fail_count += 1 - - progress.advance(upgrade_task) + # Process results as they complete + for future in as_completed(future_to_pkg): + pkg_name, latest_version, success = future.result() + + if success: + progress.console.print(f" ✅ [green]Successfully upgraded {pkg_name} to {latest_version}[/green]") + success_count += 1 + else: + progress.console.print(f" ❌ [red]Failed to upgrade {pkg_name}[/red]") + fail_count += 1 + + # Advance the progress bar for each completed task + progress.advance(upgrade_task) - # --- Summary Report --- + # --- Summary Report (Unchanged) --- console.print("\n--- [bold]Upgrade Complete[/bold] ---") console.print(f"[green]Successfully upgraded:[/green] {success_count} packages") if fail_count > 0: console.print(f"[red]Failed to upgrade:[/red] {fail_count} packages") console.print("--------------------------") - -# This makes the script runnable directly, though it's meant to be installed via the entry point if __name__ == "__main__": app() \ No newline at end of file