Python progress bar

progressbar.py

class ProgressBar:
    """
    A progress bar that's heavily inspired by click.ProgressBar.


    Required imports:
        import shutil
        import sys
        import time
        import unicodedata
        from typing import IO, Iterable

    Examples:
        data = range(100)
        p = ProgressBar(total=100, iterator=data)
        for item in p:
            print(item)
            time.sleep(1)  # do a thing

        data = range(10)
        for item in ProgressBar(total=100, iterator=data):
            print(item)
            time.sleep(1)

        data = range(10)
        p = ProgressBar(total=10)
        for index, item in enumerate(data):
            p.update(index)
            time.sleep(1)  # do a thing
    """

    _cls = "\033[1;J"

    def __init__(
        self,
        total: int,
        line_width: int | None = None,
        file: IO[str] = sys.stdout,
        fill_char: str = "#",
        empty_char: str = " ",
        iterable: Iterable | None = None,
        bar_starting_char: str = "-[",
        bar_ending_char: str = "]-",
    ):
        self.bar_ending_char = bar_ending_char
        self.bar_starting_char = bar_starting_char
        self.empty_char = empty_char
        self.file = file
        self.fill_char = fill_char
        self.iterable = iterable
        self.total = total
        self.line_width = line_width

        self.chunk = 0
        self.eta = None
        self.finished: bool = False
        self.percentage = 0
        self.plus_one: bool = False
        self.t_end = None
        self.t_start = time.time()

        size = shutil.get_terminal_size()
        self.term_columns = size.columns or 80
        self.term_rows = size.lines or 24

        self.message = None

        if not iterable:
            self.iterable = range(total)

        if len(self.fill_char) != 1:
            raise ValueError("fill_char must be a single character")

        if not self.file.isatty():
            print("WARNING: your file is not a tty; this is probably not what you want")

    def __iter__(self):
        for index, item in enumerate(self.iterable):
            yield item
            self.update(index)

        self.finished = True
        self.t_end = time.time()

    def _format_timespan(self, seconds: int | None = None) -> str:
        seconds = seconds or self.eta or 0
        hours, seconds = divmod(seconds, 60 * 60)
        minutes, seconds = divmod(seconds, 60)
        return f"{hours:02.0f}h:{minutes:02.0f}m:{seconds:02.0f}s"

    def write(self, message: str):
        self.message = message
        self.display()

    def clear_screen(self):
        self.file.write(self._cls)
        self.file.flush()

    def update(self, chunk: int, display: bool = True):
        # Make this simpler to use.  If we get passed `update(0)`, assume the
        # user _actually_ wanted to pass in `i + 1`.
        if not chunk:
            self.plus_one = True
            chunk = 1
        elif self.plus_one:
            chunk += 1

        self.chunk = chunk
        t_passed = time.time() - self.t_start

        self.eta = (t_passed / chunk) * (self.total - chunk)
        self.finished = chunk >= self.total

        if self.finished:
            self.percentage = 1.0
            self.t_end = time.time()
        else:
            self.percentage = chunk / self.total

        if display:
            self.display()

    def _get_char_width(self, char: str) -> int:
        """A pretty naive way of figuring out relative width for a character"""
        width_map = {"A": 1, "F": 1, "H": 0.5, "N": 1, "Na": 1, "W": 2}
        result = unicodedata.east_asian_width(char)
        return width_map.get(result, 1)

    def formatted_total_time(self) -> str:
        """Returns a formatted string with the total time"""
        if self.t_end and self.t_start:
            return self._format_timespan(seconds=self.t_end - self.t_start)

        return "???"

    def display(self):
        number_of_chars_in_total = len(str(self.total))
        formatted_chunk = f"%0{number_of_chars_in_total}d" % self.chunk
        bar_footer = f" ({formatted_chunk}/{self.total}) eta: {self._format_timespan()}"
        start_end_length = len(self.bar_starting_char + self.bar_ending_char)
        bar_inner_width = self.term_columns - len(bar_footer) - start_end_length
        if self.line_width:
            bar_inner_width = self.line_width - len(bar_footer) - start_end_length

        # Calculate the actual width based on how "wide" the fill character is
        number_of_filled_chars = int(self.percentage * bar_inner_width)
        filled_text = self.fill_char * int(
            number_of_filled_chars / self._get_char_width(self.fill_char)
        )
        empty_text = self.empty_char * (bar_inner_width - number_of_filled_chars)

        # Finally start building the output
        bar = f"{self.bar_starting_char}{filled_text}{empty_text}"

        # pin the appending text at a specific position to account for wide
        # fill characters.  Otherwise the appending text will bounce with wide chars.
        footer_pos = bar_inner_width + len(self.bar_starting_char) + 1
        bar += f"\033[{footer_pos};G"
        bar += f"{self.bar_ending_char}{bar_footer}"

        # Move to the bottom
        bottom_left = f"\033[{self.term_rows};0;H"
        self.file.write(bottom_left)

        if self.message:
            self.file.write("\033[2;0;H")
            self.file.write(self.message)
            self.file.write(bottom_left)

        # clear the line and then output the bar
        self.file.write("\033[0;K")
        self.file.write(bar)
        self.file.flush()