[ Web Proxy ]
URL:
Viewing: https://raw.githubusercontent.com/feast-dev/feast/v0.63.0/sdk/python/feast/diff/apply_progress.py [Back]  [Original]

"""
Enhanced progress tracking infrastructure for feast apply operations.

This module provides the ApplyProgressContext class that manages positioned,
color-coded progress bars during apply operations with fixed-width formatting
for perfect alignment.
"""

from dataclasses import dataclass
from typing import Optional

from tqdm import tqdm

try:
    from feast.diff.progress_utils import (
        create_positioned_tqdm,
        get_color_for_phase,
        is_tty_available,
    )

    _PROGRESS_UTILS_AVAILABLE = True
except ImportError:
    # Graceful fallback when progress_utils is not available (e.g., in tests)
    _PROGRESS_UTILS_AVAILABLE = False

    def create_positioned_tqdm(
        position: int,
        description: str,
        total: int,
        color: str = "blue",
        postfix: Optional[str] = None,
    ) -> Optional[tqdm]:
        return None

    def get_color_for_phase(phase: str) -> str:
        return "blue"

    def is_tty_available() -> bool:
        return False


@dataclass
class ApplyProgressContext:
    """
    Enhanced context object for tracking progress during feast apply operations.

    This class manages multiple positioned progress bars with fixed-width formatting:
    1. Overall progress (position 0) - tracks main phases
    2. Phase progress (position 1) - tracks operations within current phase

    Features:
    - Fixed-width alignment for perfect visual consistency
    - Color-coded progress bars by phase
    - Position coordination to prevent overlap
    - TTY detection for CI/CD compatibility
    """

    # Core tracking state
    current_phase: str = ""
    overall_progress: Optional[tqdm] = None
    phase_progress: Optional[tqdm] = None

    # Progress tracking
    total_phases: int = 3
    completed_phases: int = 0
    tty_available: bool = True

    # Position allocation
    OVERALL_POSITION = 0
    PHASE_POSITION = 1

    def __post_init__(self):
        """Initialize TTY detection after dataclass creation."""
        self.tty_available = _PROGRESS_UTILS_AVAILABLE and is_tty_available()

    def start_overall_progress(self):
        """Initialize the overall progress bar for apply phases."""
        if not self.tty_available:
            return

        if self.overall_progress is None:
            try:
                self.overall_progress = create_positioned_tqdm(
                    position=self.OVERALL_POSITION,
                    description="Applying changes",
                    total=self.total_phases,
                    color=get_color_for_phase("overall"),
                )
            except (TypeError, AttributeError):
                # Handle case where fallback functions don't work as expected
                self.overall_progress = None

    def start_phase(self, phase_name: str, operations_count: int = 0):
        """
        Start tracking a new phase.

        Args:
            phase_name: Human-readable name of the phase
            operations_count: Number of operations in this phase (0 for unknown)
        """
        if not self.tty_available:
            return

        self.current_phase = phase_name

        # Close previous phase progress if exists
        if self.phase_progress:
            try:
                self.phase_progress.close()
            except (AttributeError, TypeError):
                pass
            self.phase_progress = None

        # Create new phase progress bar if operations are known
        if operations_count > 0:
            try:
                self.phase_progress = create_positioned_tqdm(
                    position=self.PHASE_POSITION,
                    description=phase_name,
                    total=operations_count,
                    color=get_color_for_phase(phase_name.lower()),
                )
            except (TypeError, AttributeError):
                # Handle case where fallback functions don't work as expected
                self.phase_progress = None

    def update_phase_progress(self, description: Optional[str] = None):
        """
        Update progress within the current phase.

        Args:
            description: Optional description of current operation
        """
        if not self.tty_available or not self.phase_progress:
            return

        try:
            if description:
                # Update postfix with current operation
                self.phase_progress.set_postfix_str(description)

            self.phase_progress.update(1)
        except (AttributeError, TypeError):
            # Handle case where phase_progress is None or fallback function returned None
            pass

    def complete_phase(self):
        """Mark current phase as complete and advance overall progress."""
        if not self.tty_available:
            return

        # Close phase progress
        if self.phase_progress:
            try:
                self.phase_progress.close()
            except (AttributeError, TypeError):
                pass
            self.phase_progress = None

        # Update overall progress
        if self.overall_progress:
            try:
                self.overall_progress.update(1)
                # Update postfix with phase completion
                phase_text = f"({self.completed_phases + 1}/{self.total_phases} phases)"
                self.overall_progress.set_postfix_str(phase_text)
            except (AttributeError, TypeError):
                pass

        self.completed_phases += 1

    def cleanup(self):
        """Clean up all progress bars. Should be called in finally blocks."""
        if self.phase_progress:
            try:
                self.phase_progress.close()
            except (AttributeError, TypeError):
                pass
            self.phase_progress = None
        if self.overall_progress:
            try:
                self.overall_progress.close()
            except (AttributeError, TypeError):
                pass
            self.overall_progress = None

Web Proxy Viewer  |  New URL  |  Original Page