cdo-toolkit 0.1.0__py3-none-any.whl

cdo_toolkit/ui.py

@@ -0,0 +1,650 @@
+#!/usr/bin/env python3
+"""
+Rich UI components for CDO regridding pipeline.
+
+This module provides comprehensive progress tracking and status display
+for regridding operations, similar to the download UI system.
+"""
+
+import time
+import logging
+from datetime import datetime
+from pathlib import Path
+from typing import Optional, Dict, List, Any
+
+from rich.progress import (
+    Progress,
+    BarColumn,
+    TimeElapsedColumn,
+    TimeRemainingColumn,
+    TextColumn,
+    SpinnerColumn,
+    TaskProgressColumn,
+    MofNCompleteColumn,
+)
+from rich.console import Console
+from rich.table import Table
+from rich.panel import Panel
+
+from cdo_toolkit.errors import default_log_dir
+
+_PHASE_LABELS = {
+    "starting": "Starting",
+    "creating": "Creating chunks",
+    "regridding": "Regridding chunks",
+    "combining": "Combining chunks",
+    "prep": "Preparing",
+}
+
+
+class RegridProgressUI:
+    """Progress UI for regridding: one progress bar per file (like download UI)."""
+
+    def __init__(
+        self,
+        files: List[Path],
+        verbose: bool = True,
+        verbose_diagnostics: bool = False,
+        log_file: Optional[Path] = None,
+    ):
+        self.files = files
+        self.verbose = verbose
+        self.verbose_diagnostics = verbose_diagnostics
+        self.console = Console()
+        self.status_counts = {"completed": 0, "skipped": 0, "failed": 0}
+        self.file_status: Dict[Path, str] = {}
+        self.file_task_ids: Dict[Path, Optional[int]] = {}
+        self.failed_files: List[tuple[Path, str]] = []
+        self.processing_stats: Dict[str, Any] = {
+            "weights_reused": 0,
+            "weights_generated": 0,
+            "chunks_processed": 0,
+            "total_size_gb": 0.0,
+            "memory_peak_gb": 0.0,
+        }
+        self.overall_task: Optional[int] = None
+        self.current_files: List[Path] = []
+        self._setup_logger(log_file)
+        self._setup_progress()
+
+    def _setup_logger(self, log_file: Optional[Path] = None):
+        """Set up logging for regridding errors. Writes to log_file if given, else creates one under default_log_dir()."""
+        self.logger = logging.getLogger(f"cdo_toolkit.ui.{id(self)}")
+        self.logger.setLevel(logging.ERROR)
+
+        if log_file is None:
+            log_dir = default_log_dir()
+            log_dir.mkdir(parents=True, exist_ok=True)
+            timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+            log_file = log_dir / f"regrid_errors_{timestamp}.log"
+        self.log_file = Path(log_file)
+        self.log_file.parent.mkdir(parents=True, exist_ok=True)
+
+        handler = logging.FileHandler(self.log_file)
+        formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
+        handler.setFormatter(formatter)
+        if self.logger.hasHandlers():
+            self.logger.handlers.clear()
+        self.logger.addHandler(handler)
+
+    def _setup_progress(self):
+        """Set up compact progress display components."""
+        self.progress = Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(),
+            TaskProgressColumn(),
+            MofNCompleteColumn(),
+            TimeElapsedColumn(),
+            TimeRemainingColumn(),
+            transient=False,
+            expand=True,
+            auto_refresh=True,
+            refresh_per_second=2,  # Reduced refresh rate
+        )
+
+    def __enter__(self):
+        """Enter the progress context. Only overall task is added; per-file tasks added when each file starts."""
+        self.progress.__enter__()
+        file_str = "files" if len(self.files) > 1 else "file"
+        self.overall_task = self.progress.add_task(
+            f"[cyan]Regridding {len(self.files)} {file_str}[/cyan]",
+            total=len(self.files),
+        )
+        for file_path in self.files:
+            self.file_status[file_path] = "PENDING"
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit the progress context."""
+        self.progress.__exit__(exc_type, exc_val, exc_tb)
+
+    def start_file_processing(
+        self,
+        file_path: Path,
+        file_info: Dict[str, Any],
+        regrid_mode: str = "complete",
+    ):
+        """Add a progress bar for this file (only current batch visible). regrid_mode: 'surface', 'seafloor', or 'complete'."""
+        self.current_files.append(file_path)
+        self.file_status[file_path] = "PROCESSING"
+        filename = file_path.name
+        mode_str = f" ({regrid_mode})"
+        if self.verbose_diagnostics:
+            size_gb = file_info.get("file_size_gb", 0)
+            grid_type = file_info.get("grid_type", "unknown")
+            desc = f"[blue]{filename}{mode_str}[/blue] ({size_gb:.1f}GB, {grid_type})"
+        else:
+            desc = f"[blue][PENDING]{mode_str} {filename}[/blue]"
+        task_id = self.progress.add_task(desc, total=100, visible=True)
+        self.file_task_ids[file_path] = task_id
+
+    def update_file_progress(
+        self, file_path: Path, progress: int, operation: str = "", regrid_mode: str = "complete"
+    ):
+        """Update progress for a specific file."""
+        task_id = self.file_task_ids.get(file_path)
+        if task_id is None:
+            return
+        filename = file_path.name
+        mode_str = f" ({regrid_mode})"
+        desc = f"[blue]{operation}: {filename}{mode_str}[/blue]" if operation else f"[blue]{filename}{mode_str}[/blue]"
+        self.progress.update(task_id, completed=progress, description=desc)
+
+    def complete_file(self, file_path: Path, success: bool = True, message: str = ""):
+        """Mark a file as completed with compact display."""
+        # Remove from current processing
+        if file_path in self.current_files:
+            self.current_files.remove(file_path)
+        
+        # Update status
+        if success:
+            self.file_status[file_path] = "COMPLETED"
+            self.status_counts["completed"] += 1
+        else:
+            self.file_status[file_path] = "FAILED"
+            self.status_counts["failed"] += 1
+            if message:
+                self.failed_files.append((file_path, message))
+                self.logger.error(f"File: {file_path.name} - {message}")
+        
+        task_id = self.file_task_ids.get(file_path)
+        if task_id is not None:
+            filename = file_path.name
+            if success:
+                self.progress.update(task_id, completed=100, description=f"[green]✓ {filename}[/green]")
+            else:
+                self.progress.update(task_id, completed=100, description=f"[red]✗ {filename}[/red]")
+            self._hide_task_after_delay(task_id, delay_seconds=2)
+        # Advance overall progress
+        if self.overall_task is not None:
+            self.progress.advance(self.overall_task)
+            
+        # Update overall progress description with current stats
+        self._update_overall_description()
+
+    def skip_file(self, file_path: Path, reason: str = "Already exists"):
+        """Mark a file as skipped."""
+        # Remove from current processing
+        if file_path in self.current_files:
+            self.current_files.remove(file_path)
+        
+        self.file_status[file_path] = "SKIPPED"
+        self.status_counts["skipped"] += 1
+        
+        task_id = self.file_task_ids.get(file_path)
+        if task_id is not None:
+            filename = file_path.name
+            self.progress.update(task_id, completed=100, description=f"[yellow]⏭ {filename}[/yellow]")
+            self._hide_task_after_delay(task_id, delay_seconds=2)
+        # Advance overall progress
+        if self.overall_task is not None:
+            self.progress.advance(self.overall_task)
+            
+        # Update overall progress description
+        self._update_overall_description()
+
+    def _hide_task_after_delay(self, task_id: int, delay_seconds: int = 2):
+        """Hide a task after a delay so only the current batch remains visible."""
+        import threading
+        def hide():
+            time.sleep(delay_seconds)
+            try:
+                self.progress.update(task_id, visible=False)
+            except Exception:
+                pass
+        threading.Thread(target=hide, daemon=True).start()
+
+    def _update_overall_description(self):
+        """Update the overall progress description with current statistics."""
+        if self.overall_task is not None:
+            total = len(self.files)
+            completed = self.status_counts["completed"]
+            skipped = self.status_counts["skipped"]
+            failed = self.status_counts["failed"]
+            processing = len(self.current_files)
+            
+            desc = f"[cyan]Regridding {total} files[/cyan] - "
+            desc += f"[green]✓{completed}[/green] "
+            desc += f"[yellow]⏭{skipped}[/yellow] "
+            desc += f"[red]✗{failed}[/red]"
+            if processing > 0:
+                desc += f" [blue]Processing {processing}[/blue]"
+            
+            self.progress.update(self.overall_task, description=desc)
+
+    def update_chunk_progress(
+        self,
+        file_path: Path,
+        chunks_done: int,
+        total_chunks: int,
+        phase: str = "creating",
+        regrid_mode: str = "complete",
+    ):
+        """Update per-file progress from chunk counts (creation or regridding phase)."""
+        task_id = self.file_task_ids.get(file_path)
+        if task_id is None:
+            return
+        label = _PHASE_LABELS.get(phase, phase)
+        filename = file_path.name
+        mode_str = f" ({regrid_mode})" if regrid_mode != "complete" else ""
+        total = max(total_chunks, 1)
+        count_str = f" ({chunks_done}/{total_chunks})" if total_chunks > 0 else ""
+        self.progress.update(
+            task_id,
+            total=total,
+            completed=min(chunks_done, total),
+            description=f"[cyan]{label} {filename}{mode_str}{count_str}[/cyan]",
+        )
+
+    def update_chunking_progress(self, file_path: Path, chunk_num: int, total_chunks: int):
+        """Backward-compatible alias for chunk-creation progress."""
+        self.update_chunk_progress(file_path, chunk_num, total_chunks, phase="creating")
+
+    def update_regridding_progress(self, file_path: Path, operation: str):
+        """Update progress for regridding operations."""
+        task_id = self.file_task_ids.get(file_path)
+        if task_id is not None:
+            filename = file_path.name
+            self.progress.update(task_id, description=f"[magenta]{operation}: {filename}[/magenta]")
+
+    def _update_stats(self, stats: Dict[str, Any]):
+        """Update processing statistics."""
+        self.processing_stats.update(stats)
+
+
+    def print_summary(self):
+        """Print a comprehensive summary of the regridding operation."""
+        # Stop the progress display
+        self.progress.stop()
+        
+        # Create summary table
+        table = Table(show_header=True, header_style="bold magenta")
+        table.add_column("Status", style="bold")
+        table.add_column("Count", style="bold")
+        table.add_column("Percentage", style="bold")
+        
+        total_files = len(self.files)
+        completed = self.status_counts["completed"]
+        skipped = self.status_counts["skipped"]
+        failed = self.status_counts["failed"]
+        
+        table.add_row("[green]Completed[/green]", str(completed), f"{(completed/total_files*100):.1f}%")
+        table.add_row("[yellow]Skipped[/yellow]", str(skipped), f"{(skipped/total_files*100):.1f}%")
+        table.add_row("[red]Failed[/red]", str(failed), f"{(failed/total_files*100):.1f}%")
+        
+        # Add processing statistics
+        stats_table = Table(show_header=True, header_style="bold cyan")
+        stats_table.add_column("Metric", style="bold")
+        stats_table.add_column("Value", style="bold")
+        
+        stats_table.add_row("Weights Reused", str(self.processing_stats["weights_reused"]))
+        stats_table.add_row("Weights Generated", str(self.processing_stats["weights_generated"]))
+        stats_table.add_row("Chunks Processed", str(self.processing_stats["chunks_processed"]))
+        stats_table.add_row("Total Size (GB)", f"{self.processing_stats['total_size_gb']:.2f}")
+        stats_table.add_row("Peak Memory (GB)", f"{self.processing_stats['memory_peak_gb']:.2f}")
+        
+        # Add timing information if available
+        if "processing_time" in self.processing_stats:
+            stats_table.add_row("Processing Time", self.processing_stats["processing_time"])
+        
+        # Display summary
+        self.console.print(Panel(table, title="[bold]Regridding Summary[/bold]", border_style="green"))
+        self.console.print(Panel(stats_table, title="[bold]Processing Statistics[/bold]", border_style="cyan"))
+        
+        # Show failed files if any
+        if self.failed_files:
+            self._display_failed_files()
+
+    def _display_failed_files(self):
+        """Display information about failed files."""
+        self.console.print(f"\n[red]Failed Files ({len(self.failed_files)}):[/red]")
+        
+        for i, (file_path, error_msg) in enumerate(self.failed_files):
+            self.console.print(f"[red]{i+1}. {file_path.name}[/red]")
+            self.console.print(f"   [red]Error: {error_msg}[/red]")
+            self.console.print(f"   [red]Path: {file_path}[/red]")
+            self.console.print("")
+        
+        # Show detailed failure info for first failed file
+        if self.failed_files:
+            file_path, error_msg = self.failed_files[0]
+            self.console.print(
+                Panel(
+                    f"First failed file: [bold]{file_path.name}[/bold]\n[red]{error_msg}[/red]\n\nPath: {file_path}",
+                    title="[red]Detailed Failure Info[/red]",
+                    style="red",
+                )
+            )
+
+
+class BatchRegridUI:
+    """Batch regridding UI: only files in the current batch shown at once; overall progress."""
+
+    def __init__(
+        self,
+        files: List[Path],
+        max_workers: int = 4,
+        verbose: bool = True,
+        regrid_mode: str = "complete",
+    ):
+        self.files = files
+        self.max_workers = max_workers
+        self.verbose = verbose
+        self.regrid_mode = regrid_mode
+        self.console = Console()
+        self.overall_progress: Optional[int] = None
+        self.completed_files: List[Path] = []
+        self.failed_files: List[tuple[Path, str]] = []
+        self.skipped_files: List[Path] = []
+        self.stats = {
+            "files_processed": 0,
+            "weights_reused": 0,
+            "weights_generated": 0,
+            "chunks_processed": 0,
+            "errors": 0,
+            "total_size_gb": 0.0,
+            "memory_peak_gb": 0.0,
+            "processing_time": "0s",
+        }
+        self.file_task_ids: Dict[str, int] = {}
+        self.active_files: set[str] = set()
+        self.finished_paths: set[str] = set()
+        self._setup_progress()
+
+    def _setup_progress(self):
+        """Set up progress display: overall + one bar per file."""
+        self.progress = Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(),
+            TaskProgressColumn(),
+            MofNCompleteColumn(),
+            TimeElapsedColumn(),
+            TimeRemainingColumn(),
+            transient=False,
+            expand=True,
+            auto_refresh=True,
+            refresh_per_second=2,
+        )
+
+    def __enter__(self):
+        """Enter the batch progress context. Per-file tasks added only when each completes (so only current batch visible)."""
+        self.progress.__enter__()
+        mode_str = f" ({self.regrid_mode})"
+        desc = f"[cyan]Batch Regridding {len(self.files)} files{mode_str} ({self.max_workers} workers)[/cyan]"
+        self.overall_progress = self.progress.add_task(desc, total=len(self.files))
+        return self
+
+    def _hide_task_after_delay(self, task_id: int, delay_seconds: int = 2):
+        """Hide a task after a delay so only the current batch remains visible."""
+        import threading
+        def hide():
+            time.sleep(delay_seconds)
+            try:
+                self.progress.update(task_id, visible=False)
+            except Exception:
+                pass
+        threading.Thread(target=hide, daemon=True).start()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit the batch progress context."""
+        self.progress.__exit__(exc_type, exc_val, exc_tb)
+
+    def start_file_processing(self, file_path: Path, total_chunks: int = 0):
+        """Show a progress bar when a worker starts processing this file."""
+        path_key = str(file_path)
+        if path_key in self.file_task_ids or path_key in self.finished_paths:
+            return
+        self.active_files.add(path_key)
+        mode_str = f" ({self.regrid_mode})" if self.regrid_mode != "complete" else ""
+        total = max(total_chunks, 1) if total_chunks > 0 else 100
+        desc = f"[blue]Starting {file_path.name}{mode_str}[/blue]"
+        task_id = self.progress.add_task(desc, total=total, completed=0, visible=True)
+        self.file_task_ids[path_key] = task_id
+        self._update_overall_description()
+
+    def update_chunk_progress(
+        self,
+        file_path: Path,
+        chunks_done: int,
+        total_chunks: int,
+        phase: str = "creating",
+    ):
+        """Update in-progress bar from chunk counts reported by a worker."""
+        path_key = str(file_path)
+        if path_key not in self.file_task_ids:
+            self.start_file_processing(file_path, total_chunks=total_chunks)
+        task_id = self.file_task_ids.get(path_key)
+        if task_id is None:
+            return
+        label = _PHASE_LABELS.get(phase, phase)
+        mode_str = f" ({self.regrid_mode})" if self.regrid_mode != "complete" else ""
+        total = max(total_chunks, 1) if total_chunks > 0 else 100
+        count_str = f" ({chunks_done}/{total_chunks})" if total_chunks > 0 else ""
+        self.progress.update(
+            task_id,
+            total=total,
+            completed=min(chunks_done, total) if total_chunks > 0 else chunks_done,
+            description=f"[cyan]{label} {file_path.name}{mode_str}{count_str}[/cyan]",
+        )
+
+    def sync_from_progress_state(self, progress_state: Dict[str, Dict[str, Any]]):
+        """Poll shared worker progress and refresh in-progress bars."""
+        for path_str, info in progress_state.items():
+            if path_str in self.finished_paths or info.get("phase") == "done":
+                continue
+            file_path = Path(path_str)
+            chunks_done = int(info.get("chunks_done", 0))
+            chunks_total = int(info.get("chunks_total", 0))
+            phase = info.get("phase", "starting")
+            self.update_chunk_progress(file_path, chunks_done, chunks_total, phase=phase)
+        self._update_overall_description()
+
+    def _complete_in_progress_file(self, file_path: Path, result: Dict[str, Any]) -> Optional[int]:
+        """Mark an in-progress file complete and schedule hiding its bar."""
+        path_key = str(file_path)
+        self.finished_paths.add(path_key)
+        self.active_files.discard(path_key)
+        task_id = self.file_task_ids.pop(path_key, None)
+        if task_id is None:
+            return None
+        name = file_path.name
+        mode_str = f" ({self.regrid_mode})" if self.regrid_mode != "complete" else ""
+        if result.get("success", False):
+            if result.get("skipped", False):
+                desc = f"[yellow]⏭ {name}{mode_str}[/yellow]"
+            else:
+                desc = f"[green]✓ {name}{mode_str}[/green]"
+        else:
+            desc = f"[red]✗ {name}{mode_str}[/red]"
+        self.progress.update(task_id, total=1, completed=1, description=desc)
+        self._hide_task_after_delay(task_id, delay_seconds=2)
+        return task_id
+
+    def update_file_result(self, file_path: Path, result: Dict[str, Any]):
+        """Update progress when a worker finishes; complete in-progress bar or flash result briefly."""
+        name = file_path.name
+        mode_str = f" ({self.regrid_mode})" if self.regrid_mode != "complete" else ""
+        if result.get("success", False):
+            if result.get("skipped", False):
+                self.skipped_files.append(file_path)
+                desc = f"[yellow]⏭ {name}{mode_str}[/yellow]"
+            else:
+                self.completed_files.append(file_path)
+                desc = f"[green]✓ {name}{mode_str}[/green]"
+        else:
+            error_msg = result.get("message", "Unknown error")
+            self.failed_files.append((file_path, error_msg))
+            desc = f"[red]✗ {name}{mode_str}[/red]"
+
+        if self._complete_in_progress_file(file_path, result) is None:
+            task_id = self.progress.add_task(desc, total=1, completed=1, visible=True)
+            self._hide_task_after_delay(task_id, delay_seconds=2)
+
+        if self.overall_progress is not None:
+            self.progress.advance(self.overall_progress)
+        self._update_overall_description()
+
+        if "stats" in result:
+            self._update_stats(result["stats"])
+
+    def _update_overall_description(self):
+        """Update the overall progress description with current statistics."""
+        if self.overall_progress is not None:
+            total = len(self.files)
+            completed = len(self.completed_files)
+            skipped = len(self.skipped_files)
+            failed = len(self.failed_files)
+            remaining = total - completed - skipped - failed
+            
+            in_progress = len(self.active_files)
+            desc = f"[cyan]Batch Regridding {total} files ({self.max_workers} workers)[/cyan] - "
+            desc += f"[green]✓{completed}[/green] "
+            desc += f"[yellow]⏭{skipped}[/yellow] "
+            desc += f"[red]✗{failed}[/red]"
+            if in_progress > 0:
+                desc += f" [blue]In progress: {in_progress}[/blue]"
+            elif remaining > 0:
+                desc += f" [blue]Remaining: {remaining}[/blue]"
+            
+            self.progress.update(self.overall_progress, description=desc)
+
+    def _update_stats(self, new_stats: Dict[str, Any]):
+        """Update cumulative statistics."""
+        for key, value in new_stats.items():
+            if key in self.stats:
+                if isinstance(value, (int, float)):
+                    self.stats[key] += value
+                else:
+                    self.stats[key] = value
+            else:
+                self.stats[key] = value # add new key to stats dictionary
+
+    def print_summary(self):
+        """Print compact batch processing summary."""
+        self.progress.stop()
+        
+        # Create compact summary
+        table = Table(show_header=True, header_style="bold magenta")
+        table.add_column("Status", style="bold")
+        table.add_column("Count", style="bold")
+        table.add_column("Percentage", style="bold")
+        
+        total = len(self.files)
+        completed = len(self.completed_files)
+        skipped = len(self.skipped_files)
+        failed = len(self.failed_files)
+        
+        table.add_row("[green]Completed[/green]", str(completed), f"{(completed/total*100):.1f}%")
+        table.add_row("[yellow]Skipped[/yellow]", str(skipped), f"{(skipped/total*100):.1f}%")
+        table.add_row("[red]Failed[/red]", str(failed), f"{(failed/total*100):.1f}%")
+        
+        # Add processing stats
+        stats_table = Table(show_header=True, header_style="bold cyan")
+        stats_table.add_column("Metric", style="bold")
+        stats_table.add_column("Value", style="bold")
+        
+        stats_table.add_row("Weights Reused", str(self.stats["weights_reused"]))
+        stats_table.add_row("Weights Generated", str(self.stats["weights_generated"]))
+        stats_table.add_row("Chunks Processed", str(self.stats["chunks_processed"]))
+        stats_table.add_row("Total Size (GB)", f"{self.stats['total_size_gb']:.2f}")
+        stats_table.add_row("Peak Memory (GB)", f"{self.stats['memory_peak_gb']:.2f}")
+        
+        # Add timing information if available
+        if "processing_time" in self.stats:
+            stats_table.add_row("Processing Time", self.stats["processing_time"])
+        else:
+            print("No processing time available")
+            print(self.stats)
+        
+        self.console.print(Panel(table, title="[bold]Batch Regridding Summary[/bold]", border_style="green"))
+        self.console.print(Panel(stats_table, title="[bold]Processing Statistics[/bold]", border_style="cyan"))
+        
+        # Show failed files if any (compact)
+        if self.failed_files:
+            self.console.print(f"\n[red]Failed Files ({len(self.failed_files)}):[/red]")
+            for i, (file_path, error_msg) in enumerate(self.failed_files[:5]):  # Show only first 5
+                self.console.print(f"[red]{i+1}. {file_path.name} - {error_msg}[/red]")
+            if len(self.failed_files) > 5:
+                self.console.print(f"[red]... and {len(self.failed_files) - 5} more[/red]")
+
+
+class ChunkingProgressUI:
+    """UI for chunked file processing operations."""
+
+    def __init__(self, file_path: Path, total_chunks: int, verbose: bool = True):
+        self.file_path = file_path
+        self.total_chunks = total_chunks
+        self.verbose = verbose
+        self.console = Console()
+        
+        self._setup_progress()
+
+    def _setup_progress(self):
+        """Set up progress display for chunking operations."""
+        self.progress = Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(),
+            TaskProgressColumn(),
+            MofNCompleteColumn(),
+            TimeElapsedColumn(),
+            TimeRemainingColumn(),
+            transient=True,
+            expand=True,
+            auto_refresh=True,
+            refresh_per_second=4,
+        )
+
+    def __enter__(self):
+        """Enter the chunking progress context."""
+        self.progress.__enter__()
+        
+        # Add chunking progress task
+        self.chunking_task = self.progress.add_task(
+            f"[cyan]Chunking: {self.file_path.name}[/cyan]",
+            total=self.total_chunks
+        )
+        
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit the chunking progress context."""
+        self.progress.__exit__(exc_type, exc_val, exc_tb)
+
+    def update_chunk_progress(self, chunk_num: int, operation: str = "Processing"):
+        """Update progress for chunk processing."""
+        self.progress.update(
+            self.chunking_task,
+            completed=chunk_num,
+            description=f"[cyan]{operation} chunk {chunk_num}/{self.total_chunks}: {self.file_path.name}[/cyan]"
+        )
+
+    def complete_chunking(self):
+        """Mark chunking as complete."""
+        self.progress.update(
+            self.chunking_task,
+            completed=self.total_chunks,
+            description=f"[green]✓ Chunking complete: {self.file_path.name}[/green]"
+        )