diff --git a/freqtrade/util/__init__.py b/freqtrade/util/__init__.py index f3674c194..3d227880e 100644 --- a/freqtrade/util/__init__.py +++ b/freqtrade/util/__init__.py @@ -15,7 +15,11 @@ from freqtrade.util.formatters import decimals_per_coin, fmt_coin, fmt_coin2, ro from freqtrade.util.ft_precise import FtPrecise from freqtrade.util.measure_time import MeasureTime from freqtrade.util.periodic_cache import PeriodicCache -from freqtrade.util.progress_tracker import get_progress_tracker # noqa F401 +from freqtrade.util.progress_tracker import ( # noqa F401 + ProgressLike, + get_progress_tracker, + retrieve_progress_tracker, +) from freqtrade.util.rich_progress import CustomProgress from freqtrade.util.rich_tables import print_df_rich_table, print_rich_table from freqtrade.util.template_renderer import render_template, render_template_with_fallback # noqa diff --git a/freqtrade/util/progress_tracker.py b/freqtrade/util/progress_tracker.py index ca8f5d3e9..2e339c7fb 100644 --- a/freqtrade/util/progress_tracker.py +++ b/freqtrade/util/progress_tracker.py @@ -1,6 +1,10 @@ +from contextlib import AbstractContextManager +from typing import Optional, Protocol + from rich.progress import ( BarColumn, MofNCompleteColumn, + TaskID, TaskProgressColumn, TextColumn, TimeElapsedColumn, @@ -10,7 +14,19 @@ from rich.progress import ( from freqtrade.util.rich_progress import CustomProgress -def get_progress_tracker(**kwargs): +class ProgressLike(Protocol, AbstractContextManager["ProgressLike"]): + def add_task(self, description: str, *args, **kwargs) -> TaskID: ... + + def update(self, task_id: TaskID, *, advance: float | None = None, **kwargs): ... + + +def retrieve_progress_tracker(pt: Optional[ProgressLike]) -> ProgressLike: + if pt is None: + return get_progress_tracker() + return pt + + +def get_progress_tracker(**kwargs) -> ProgressLike: """ Get progress Bar with custom columns. """