utils.py

# ! ==================================== GLOBAL UTILS =======================================
# ! >>>> ABOUT PRINTOUTS
from rich.console import Console
from rich.progress import (
    BarColumn,
    Progress,
    SpinnerColumn,
    TextColumn,
    TimeRemainingColumn,
)
from typing import Iterable, Optional, Dict, Any

console = Console()

def printinfo(message: str, end="\n"):
    console.print(f"[bold blue][INFO][/bold blue] ---- {message}", end=end)


def printwarning(message: str, end="\n"):
    console.print(f"[bold yellow][WARN][/bold yellow] ---- {message}", end=end)


def printerr(message: str, end="\n"):
    console.print(f"[bold red][ERRO][/bold red] ---- {message}", end=end)

def printimp(message: str, end="\n"):
    console.print(f"[bold orange][IMPO][/bold orange] ---- {message}", end=end)

# 我们可以把 Progress 的列定义为常量，方便复用
DEFAULT_RICH_COLUMNS = [
    SpinnerColumn(),
    TextColumn("[progress.description]{task.description}"),
    BarColumn(),
    TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
    TimeRemainingColumn(),
    # 添加一个用于显示后缀的列，初始为空
    TextColumn("[green]{task.fields[postfix]}", justify="right"),
]

class ProgressBar:
    def __init__(
        self,
        sequence: Iterable,
        description: str = ">>>",
        total: Optional[int] = None,
        console: Optional[Console] = None,
        transient: bool = False,
    ):
        self.sequence = sequence
        self.description = description
        self.total = total or len(sequence)
        
        self.progress = Progress(
            *DEFAULT_RICH_COLUMNS,
            console=console or Console(),
            transient=transient,
        )
        self.task_id = self.progress.add_task(self.description, total=self.total, postfix="")

    def set_postfix(self, **kwargs: Dict[str, Any]):
        postfix_str = ", ".join(
            [f"{k}={v:.4f}" if isinstance(v, float) else f"{k}={v}" for k, v in kwargs.items()]
        )
        self.progress.update(self.task_id, postfix=postfix_str)

    def __iter__(self):
        with self.progress:
            for item in self.sequence:
                yield item
                self.progress.update(self.task_id, advance=1)

    def __len__(self):
        return self.total

# ! =========================================================================================