CoolFace
Apppublic

Aluode/PerceptionLabPortable

sourceHugging Faceupdated 9mo agoView on Hugging Face
0likes
progress_bar.py224 linesDownload Raw Back to rich
1import math2from functools import lru_cache3from time import monotonic4from typing import Iterable, List, Optional5 6from .color import Color, blend_rgb7from .color_triplet import ColorTriplet8from .console import Console, ConsoleOptions, RenderResult9from .jupyter import JupyterMixin10from .measure import Measurement11from .segment import Segment12from .style import Style, StyleType13 14# Number of characters before 'pulse' animation repeats15PULSE_SIZE = 2016 17 18class ProgressBar(JupyterMixin):19    """Renders a (progress) bar. Used by rich.progress.20 21    Args:22        total (float, optional): Number of steps in the bar. Defaults to 100. Set to None to render a pulsing animation.23        completed (float, optional): Number of steps completed. Defaults to 0.24        width (int, optional): Width of the bar, or ``None`` for maximum width. Defaults to None.25        pulse (bool, optional): Enable pulse effect. Defaults to False. Will pulse if a None total was passed.26        style (StyleType, optional): Style for the bar background. Defaults to "bar.back".27        complete_style (StyleType, optional): Style for the completed bar. Defaults to "bar.complete".28        finished_style (StyleType, optional): Style for a finished bar. Defaults to "bar.finished".29        pulse_style (StyleType, optional): Style for pulsing bars. Defaults to "bar.pulse".30        animation_time (Optional[float], optional): Time in seconds to use for animation, or None to use system time.31    """32 33    def __init__(34        self,35        total: Optional[float] = 100.0,36        completed: float = 0,37        width: Optional[int] = None,38        pulse: bool = False,39        style: StyleType = "bar.back",40        complete_style: StyleType = "bar.complete",41        finished_style: StyleType = "bar.finished",42        pulse_style: StyleType = "bar.pulse",43        animation_time: Optional[float] = None,44    ):45        self.total = total46        self.completed = completed47        self.width = width48        self.pulse = pulse49        self.style = style50        self.complete_style = complete_style51        self.finished_style = finished_style52        self.pulse_style = pulse_style53        self.animation_time = animation_time54 55        self._pulse_segments: Optional[List[Segment]] = None56 57    def __repr__(self) -> str:58        return f"<Bar {self.completed!r} of {self.total!r}>"59 60    @property61    def percentage_completed(self) -> Optional[float]:62        """Calculate percentage complete."""63        if self.total is None:64            return None65        completed = (self.completed / self.total) * 100.066        completed = min(100, max(0.0, completed))67        return completed68 69    @lru_cache(maxsize=16)70    def _get_pulse_segments(71        self,72        fore_style: Style,73        back_style: Style,74        color_system: str,75        no_color: bool,76        ascii: bool = False,77    ) -> List[Segment]:78        """Get a list of segments to render a pulse animation.79 80        Returns:81            List[Segment]: A list of segments, one segment per character.82        """83        bar = "-" if ascii else "━"84        segments: List[Segment] = []85        if color_system not in ("standard", "eight_bit", "truecolor") or no_color:86            segments += [Segment(bar, fore_style)] * (PULSE_SIZE // 2)87            segments += [Segment(" " if no_color else bar, back_style)] * (88                PULSE_SIZE - (PULSE_SIZE // 2)89            )90            return segments91 92        append = segments.append93        fore_color = (94            fore_style.color.get_truecolor()95            if fore_style.color96            else ColorTriplet(255, 0, 255)97        )98        back_color = (99            back_style.color.get_truecolor()100            if back_style.color101            else ColorTriplet(0, 0, 0)102        )103        cos = math.cos104        pi = math.pi105        _Segment = Segment106        _Style = Style107        from_triplet = Color.from_triplet108 109        for index in range(PULSE_SIZE):110            position = index / PULSE_SIZE111            fade = 0.5 + cos(position * pi * 2) / 2.0112            color = blend_rgb(fore_color, back_color, cross_fade=fade)113            append(_Segment(bar, _Style(color=from_triplet(color))))114        return segments115 116    def update(self, completed: float, total: Optional[float] = None) -> None:117        """Update progress with new values.118 119        Args:120            completed (float): Number of steps completed.121            total (float, optional): Total number of steps, or ``None`` to not change. Defaults to None.122        """123        self.completed = completed124        self.total = total if total is not None else self.total125 126    def _render_pulse(127        self, console: Console, width: int, ascii: bool = False128    ) -> Iterable[Segment]:129        """Renders the pulse animation.130 131        Args:132            console (Console): Console instance.133            width (int): Width in characters of pulse animation.134 135        Returns:136            RenderResult: [description]137 138        Yields:139            Iterator[Segment]: Segments to render pulse140        """141        fore_style = console.get_style(self.pulse_style, default="white")142        back_style = console.get_style(self.style, default="black")143 144        pulse_segments = self._get_pulse_segments(145            fore_style, back_style, console.color_system, console.no_color, ascii=ascii146        )147        segment_count = len(pulse_segments)148        current_time = (149            monotonic() if self.animation_time is None else self.animation_time150        )151        segments = pulse_segments * (int(width / segment_count) + 2)152        offset = int(-current_time * 15) % segment_count153        segments = segments[offset : offset + width]154        yield from segments155 156    def __rich_console__(157        self, console: Console, options: ConsoleOptions158    ) -> RenderResult:159        width = min(self.width or options.max_width, options.max_width)160        ascii = options.legacy_windows or options.ascii_only161        should_pulse = self.pulse or self.total is None162        if should_pulse:163            yield from self._render_pulse(console, width, ascii=ascii)164            return165 166        completed: Optional[float] = (167            min(self.total, max(0, self.completed)) if self.total is not None else None168        )169 170        bar = "-" if ascii else "━"171        half_bar_right = " " if ascii else "╸"172        half_bar_left = " " if ascii else "╺"173        complete_halves = (174            int(width * 2 * completed / self.total)175            if self.total and completed is not None176            else width * 2177        )178        bar_count = complete_halves // 2179        half_bar_count = complete_halves % 2180        style = console.get_style(self.style)181        is_finished = self.total is None or self.completed >= self.total182        complete_style = console.get_style(183            self.finished_style if is_finished else self.complete_style184        )185        _Segment = Segment186        if bar_count:187            yield _Segment(bar * bar_count, complete_style)188        if half_bar_count:189            yield _Segment(half_bar_right * half_bar_count, complete_style)190 191        if not console.no_color:192            remaining_bars = width - bar_count - half_bar_count193            if remaining_bars and console.color_system is not None:194                if not half_bar_count and bar_count:195                    yield _Segment(half_bar_left, style)196                    remaining_bars -= 1197                if remaining_bars:198                    yield _Segment(bar * remaining_bars, style)199 200    def __rich_measure__(201        self, console: Console, options: ConsoleOptions202    ) -> Measurement:203        return (204            Measurement(self.width, self.width)205            if self.width is not None206            else Measurement(4, options.max_width)207        )208 209 210if __name__ == "__main__":  # pragma: no cover211    console = Console()212    bar = ProgressBar(width=50, total=100)213 214    import time215 216    console.show_cursor(False)217    for n in range(0, 101, 1):218        bar.update(n)219        console.print(bar)220        console.file.write("\r")221        time.sleep(0.05)222    console.show_cursor(True)223    console.print()224