Source code for torchmate.utils.progress_bar

  1import time
  2
  3

[docs]
  4class ProgressBar:
  5    """A progress bar for tracking the progress of iterative tasks.
  6
  7    Displays a progress bar with percentage completion, elapsed time, estimated time remaining,
  8    and optional verbose text.
  9
 10    **Example usage:**
 11
 12    .. code-block:: python
 13
 14        no_iter = 100
 15        pb = ProgressBar(total=no_iter, prefix="Processing", bar_length=50)
 16        for i in range(no_iter):
 17            pb.update(i + 1, message=f"Processing item {i + 1}")
 18
 19    """
 20
 21    def __init__(self, total: int, bar_length: int = 30, fill: str = "=", prefix: str = "") -> None:
 22        """Initialize the instance of ProgressBar class.
 23
 24        Args:
 25            total (int): Total number of iterations. Defaults to None.
 26            bar_length (int, optional): Length of the progress bar. Defaults to ``30``.
 27            fill (str, optional): Character used to fill the progress bar. Defaults to ``'='``.
 28            prefix (str, optional): A text to display before the progress bar. Defaults to empty string ``""``.
 29        """
 30
 31        self.total = total
 32        self.length = bar_length
 33        self.fill = fill
 34        self.prefix = prefix + ": " if prefix else ""
 35        self.start_time = time.time()
 36        self._prev_len = 0
 37

[docs]
 38    def update(self, iteration: int, message: str = "") -> None:
 39        """
 40        Update the progress bar with current iteration and optional message text.
 41
 42        Args:
 43            iteration (int): Current iteration number.
 44            message (str, optional): Optional text to display. Defaults to empty string.
 45
 46        """
 47
 48        # Calculate progress and format bar
 49        percent = ("{0:.1f}").format(100 * (iteration / float(self.total)))
 50        filled_length = int(self.length * iteration // self.total)
 51        bar = self.fill * filled_length + "-" * (self.length - filled_length)
 52
 53        # Calculate elapsed and estimated time
 54        elapsed_time = time.time() - self.start_time
 55        ET = self.format_time(elapsed_time)
 56        estimated_time_remaining = (elapsed_time / iteration) * (self.total - iteration)
 57        ETR = self.format_time(estimated_time_remaining)
 58
 59        # Print progress bar with appropriate text
 60        if iteration == self.total:
 61            step_time = self.format_time(elapsed_time / self.total)
 62            info = f"\r{self.prefix}{iteration}/{self.total} [{bar}] {percent}% | ET: {ET} | {step_time}\\step | "
 63            text = info + message
 64            end = "\n"
 65        else:
 66            info = f"\r{self.prefix}{iteration}/{self.total} [{bar}] {percent}% | ET: {ET}| ETR : {ETR} | "
 67            text = info + message
 68            end = "\r"
 69        print(text, end=end, flush=True)

 70

[docs]
 71    @staticmethod
 72    def format_time(duration_seconds: float) -> str:
 73        """
 74        Format a duration in seconds into a human-readable string.
 75
 76        Args:
 77            duration_seconds (float): Duration in seconds.
 78
 79        Returns:
 80            str: Formatted time string (e.g., "1m 23s", "45.67s", "2h 35m").
 81
 82        **Example usage:**
 83
 84        .. code-block:: python
 85
 86            formatted_time = ProgressBar.format_time(65.5)  # Returns "1:05"
 87
 88        """
 89
 90        if duration_seconds < 1:
 91            formated_time = f"{duration_seconds * 1e3:.0f}ms"
 92        elif duration_seconds < 60:
 93            formated_time = f"{duration_seconds:.2f}s"
 94        elif duration_seconds < 3600:
 95            minutes = int(duration_seconds // 60)
 96            seconds = int(duration_seconds % 60)
 97            formated_time = f"{minutes:02d}:{seconds:02d}min"
 98        else:
 99            hours = int(duration_seconds // 3600)
100            minutes = int((duration_seconds % 3600) // 60)
101            seconds = int(duration_seconds % 60)
102            formated_time = f"{hours:02d}:{minutes:02d}:{seconds:02d}"
103        return formated_time