diff --git a/plextraktsync/plan/Walker.py b/plextraktsync/plan/Walker.py index ef10e4eddc..9102078f1e 100644 --- a/plextraktsync/plan/Walker.py +++ b/plextraktsync/plan/Walker.py @@ -109,7 +109,7 @@ def find_episodes(self): def walk_shows(self, shows: set[Media], title="Processing Shows"): if not shows: return - yield from self.progressbar(shows, desc=title) + yield from self.progressbar(shows, desc=title, total=len(shows)) def get_plex_episodes(self, episodes: list[Episode]) -> Generator[Media, Any, None]: it = self.progressbar(episodes, desc="Processing episodes") diff --git a/plextraktsync/rich/RichProgressBar.py b/plextraktsync/rich/RichProgressBar.py new file mode 100644 index 0000000000..077ea12b95 --- /dev/null +++ b/plextraktsync/rich/RichProgressBar.py @@ -0,0 +1,64 @@ +from functools import cached_property + + +class RichProgressBar: + def __init__(self, iterable, total, options=None, desc=""): + self.iter = iter(iterable) + self.options = options or {} + self.desc = desc + self.total = total + self.i = 0 + + def __iter__(self): + return self + + def __next__(self): + res = self.iter.__next__() + self.update() + return res + + def __enter__(self): + self.progress.__enter__() + return self + + def __exit__(self, *exc): + self.progress.__exit__(*exc) + + def update(self): + self.i += 1 + self.progress.update(self.task_id, completed=self.i) + + @cached_property + def task_id(self): + return self.progress.add_task(self.desc, total=self.total) + + @cached_property + def progress(self): + from tqdm.rich import FractionColumn, RateColumn + + from rich.progress import (BarColumn, Progress, TimeElapsedColumn, + TimeRemainingColumn) + + args = ( + "[progress.description]{task.description}" + "[progress.percentage]{task.percentage:>4.0f}%", + BarColumn(bar_width=None), + FractionColumn( + unit_scale=False, + unit_divisor=1000, + ), + "[", + TimeElapsedColumn(), + "<", + TimeRemainingColumn(), + ",", + RateColumn( + unit="it", + unit_scale=False, + unit_divisor=1000, + ), + "]" + ) + progress = Progress(*args, **self.options) + + return progress diff --git a/plextraktsync/util/Factory.py b/plextraktsync/util/Factory.py index bec84bd970..4a4add0c15 100644 --- a/plextraktsync/util/Factory.py +++ b/plextraktsync/util/Factory.py @@ -140,23 +140,11 @@ def progressbar(self): if not self.run_config.progressbar: return None - import warnings from functools import partial - from tqdm import TqdmExperimentalWarning - from tqdm.rich import tqdm + from plextraktsync.rich.RichProgressBar import RichProgressBar - warnings.filterwarnings("ignore", category=TqdmExperimentalWarning) - - # Monkeypatch https://github.com/tqdm/tqdm/pull/1395 - class Tqdm(tqdm): - def close(self): - if self.disable: - return - self.display() - super().close() - - return partial(Tqdm, options={'console': self.console}) + return partial(RichProgressBar, options={'console': self.console}) @cached_property def run_config(self):