codekingpro/portable-devtools

1from dataclasses import dataclass2from typing import TYPE_CHECKING, Optional3 4from httpie.context import Environment5 6if TYPE_CHECKING:7    from rich.console import Console8 9 10@dataclass11class BaseDisplay:12    env: Environment13 14    def start(15        self, *, total: Optional[float], at: float, description: str16    ) -> None:17        ...18 19    def update(self, steps: float) -> None:20        ...21 22    def stop(self, time_spent: float) -> None:23        ...24 25    @property26    def console(self) -> 'Console':27        """Returns the default console to be used with displays (stderr)."""28        return self.env.rich_error_console29 30    def _print_summary(31        self, is_finished: bool, observed_steps: int, time_spent: float32    ):33        from rich import filesize34 35        if is_finished:36            verb = 'Done'37        else:38            verb = 'Interrupted'39 40        total_size = filesize.decimal(observed_steps)41        avg_speed = filesize.decimal(observed_steps / time_spent)42 43        minutes, seconds = divmod(time_spent, 60)44        hours, minutes = divmod(int(minutes), 60)45        if hours:46            total_time = f'{hours:d}:{minutes:02d}:{seconds:0.5f}'47        else:48            total_time = f'{minutes:02d}:{seconds:0.5f}'49 50        self.console.print(51            f'[progress.description]{verb}. {total_size} in {total_time} ({avg_speed}/s)'52        )53 54 55class DummyDisplay(BaseDisplay):56    """57    A dummy display object to be used when the progress bars,58    spinners etc. are disabled globally (or during tests).59    """60 61 62class StatusDisplay(BaseDisplay):63    def start(64        self, *, total: Optional[float], at: float, description: str65    ) -> None:66        self.observed = at67        self.description = (68            f'[progress.description]{description}[/progress.description]'69        )70 71        self.status = self.console.status(self.description, spinner='line')72        self.status.start()73 74    def update(self, steps: float) -> None:75        from rich import filesize76 77        self.observed += steps78 79        observed_amount, observed_unit = filesize.decimal(80            self.observed81        ).split()82        self.status.update(83            status=f'{self.description} [progress.download]{observed_amount}/? {observed_unit}[/progress.download]'84        )85 86    def stop(self, time_spent: float) -> None:87        self.status.stop()88        self.console.print(self.description)89        if time_spent:90            self._print_summary(91                is_finished=True,92                observed_steps=self.observed,93                time_spent=time_spent,94            )95 96 97class ProgressDisplay(BaseDisplay):98    def start(99        self, *, total: Optional[float], at: float, description: str100    ) -> None:101        from rich.progress import (102            Progress,103            BarColumn,104            DownloadColumn,105            TimeRemainingColumn,106            TransferSpeedColumn,107        )108 109        assert total is not None110        self.console.print(f'[progress.description]{description}')111        self.progress_bar = Progress(112            '[',113            BarColumn(),114            ']',115            '[progress.percentage]{task.percentage:>3.0f}%',116            '(',117            DownloadColumn(),118            ')',119            TimeRemainingColumn(),120            TransferSpeedColumn(),121            console=self.console,122            transient=True,123        )124        self.progress_bar.start()125        self.transfer_task = self.progress_bar.add_task(126            description, completed=at, total=total127        )128 129    def update(self, steps: float) -> None:130        self.progress_bar.advance(self.transfer_task, steps)131 132    def stop(self, time_spent: Optional[float]) -> None:133        self.progress_bar.stop()134 135        if time_spent:136            [task] = self.progress_bar.tasks137            self._print_summary(138                is_finished=task.finished,139                observed_steps=task.completed,140                time_spent=time_spent,141            )142